| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915 |
- import { beforeEach, describe, expect, test, vi } from "vitest";
- import { resolveEndpointPolicy } from "@/app/v1/_lib/proxy/endpoint-policy";
- const mocks = vi.hoisted(() => ({
- pickRandomProviderWithExclusion: vi.fn(),
- recordSuccess: vi.fn(),
- recordFailure: vi.fn(async () => {}),
- getCircuitState: vi.fn(() => "closed"),
- getProviderHealthInfo: vi.fn(async () => ({
- health: { failureCount: 0 },
- config: { failureThreshold: 3 },
- })),
- updateSessionBindingSmart: vi.fn(async () => ({ updated: true, reason: "test" })),
- updateSessionProvider: vi.fn(async () => {}),
- clearSessionProvider: vi.fn(async () => {}),
- isHttp2Enabled: vi.fn(async () => false),
- getPreferredProviderEndpoints: vi.fn(async () => []),
- getEndpointFilterStats: vi.fn(async () => null),
- recordEndpointSuccess: vi.fn(async () => {}),
- recordEndpointFailure: vi.fn(async () => {}),
- isVendorTypeCircuitOpen: vi.fn(async () => false),
- recordVendorTypeAllEndpointsTimeout: vi.fn(async () => {}),
- categorizeErrorAsync: vi.fn(async () => 0),
- storeSessionSpecialSettings: vi.fn(async () => {}),
- }));
- vi.mock("@/lib/logger", () => ({
- logger: {
- debug: vi.fn(),
- info: vi.fn(),
- warn: vi.fn(),
- trace: vi.fn(),
- error: vi.fn(),
- fatal: vi.fn(),
- },
- }));
- vi.mock("@/lib/config", async (importOriginal) => {
- const actual = await importOriginal<typeof import("@/lib/config")>();
- return {
- ...actual,
- isHttp2Enabled: mocks.isHttp2Enabled,
- };
- });
- vi.mock("@/lib/provider-endpoints/endpoint-selector", () => ({
- getPreferredProviderEndpoints: mocks.getPreferredProviderEndpoints,
- getEndpointFilterStats: mocks.getEndpointFilterStats,
- }));
- vi.mock("@/lib/endpoint-circuit-breaker", () => ({
- recordEndpointSuccess: mocks.recordEndpointSuccess,
- recordEndpointFailure: mocks.recordEndpointFailure,
- }));
- vi.mock("@/lib/circuit-breaker", () => ({
- getCircuitState: mocks.getCircuitState,
- getProviderHealthInfo: mocks.getProviderHealthInfo,
- recordFailure: mocks.recordFailure,
- recordSuccess: mocks.recordSuccess,
- }));
- vi.mock("@/lib/vendor-type-circuit-breaker", () => ({
- isVendorTypeCircuitOpen: mocks.isVendorTypeCircuitOpen,
- recordVendorTypeAllEndpointsTimeout: mocks.recordVendorTypeAllEndpointsTimeout,
- }));
- vi.mock("@/lib/session-manager", () => ({
- SessionManager: {
- updateSessionBindingSmart: mocks.updateSessionBindingSmart,
- updateSessionProvider: mocks.updateSessionProvider,
- clearSessionProvider: mocks.clearSessionProvider,
- storeSessionSpecialSettings: mocks.storeSessionSpecialSettings,
- },
- }));
- vi.mock("@/app/v1/_lib/proxy/provider-selector", () => ({
- ProxyProviderResolver: {
- pickRandomProviderWithExclusion: mocks.pickRandomProviderWithExclusion,
- },
- }));
- vi.mock("@/app/v1/_lib/proxy/errors", async (importOriginal) => {
- const actual = await importOriginal<typeof import("@/app/v1/_lib/proxy/errors")>();
- return {
- ...actual,
- categorizeErrorAsync: mocks.categorizeErrorAsync,
- };
- });
- import {
- ErrorCategory as ProxyErrorCategory,
- ProxyError as UpstreamProxyError,
- } from "@/app/v1/_lib/proxy/errors";
- import { ProxyForwarder } from "@/app/v1/_lib/proxy/forwarder";
- import { ProxySession } from "@/app/v1/_lib/proxy/session";
- import type { Provider } from "@/types/provider";
- type AttemptRuntime = {
- clearResponseTimeout?: () => void;
- responseController?: AbortController;
- };
- function createProvider(overrides: Partial<Provider> = {}): Provider {
- return {
- id: 1,
- name: "p1",
- url: "https://provider.example.com",
- key: "k",
- providerVendorId: null,
- isEnabled: true,
- weight: 1,
- priority: 0,
- groupPriorities: null,
- costMultiplier: 1,
- groupTag: null,
- providerType: "claude",
- preserveClientIp: false,
- modelRedirects: null,
- allowedModels: null,
- mcpPassthroughType: "none",
- mcpPassthroughUrl: null,
- limit5hUsd: null,
- limitDailyUsd: null,
- dailyResetMode: "fixed",
- dailyResetTime: "00:00",
- limitWeeklyUsd: null,
- limitMonthlyUsd: null,
- limitTotalUsd: null,
- totalCostResetAt: null,
- limitConcurrentSessions: 0,
- maxRetryAttempts: 1,
- circuitBreakerFailureThreshold: 5,
- circuitBreakerOpenDuration: 1_800_000,
- circuitBreakerHalfOpenSuccessThreshold: 2,
- proxyUrl: null,
- proxyFallbackToDirect: false,
- firstByteTimeoutStreamingMs: 100,
- streamingIdleTimeoutMs: 0,
- requestTimeoutNonStreamingMs: 0,
- websiteUrl: null,
- faviconUrl: null,
- cacheTtlPreference: null,
- context1mPreference: null,
- codexReasoningEffortPreference: null,
- codexReasoningSummaryPreference: null,
- codexTextVerbosityPreference: null,
- codexParallelToolCallsPreference: null,
- codexServiceTierPreference: null,
- anthropicMaxTokensPreference: null,
- anthropicThinkingBudgetPreference: null,
- anthropicAdaptiveThinking: null,
- geminiGoogleSearchPreference: null,
- tpm: 0,
- rpm: 0,
- rpd: 0,
- cc: 0,
- createdAt: new Date(),
- updatedAt: new Date(),
- deletedAt: null,
- ...overrides,
- };
- }
- function createSession(clientAbortSignal: AbortSignal | null = null): ProxySession {
- const headers = new Headers();
- const session = Object.create(ProxySession.prototype);
- Object.assign(session, {
- startTime: Date.now(),
- method: "POST",
- requestUrl: new URL("https://example.com/v1/messages"),
- headers,
- originalHeaders: new Headers(headers),
- headerLog: JSON.stringify(Object.fromEntries(headers.entries())),
- request: {
- model: "claude-test",
- log: "(test)",
- message: {
- model: "claude-test",
- stream: true,
- messages: [{ role: "user", content: "hi" }],
- },
- },
- userAgent: null,
- context: null,
- clientAbortSignal,
- userName: "test-user",
- authState: { success: true, user: null, key: null, apiKey: null },
- provider: null,
- messageContext: null,
- sessionId: "sess-hedge",
- requestSequence: 1,
- originalFormat: "claude",
- providerType: null,
- originalModelName: null,
- originalUrlPathname: null,
- providerChain: [],
- cacheTtlResolved: null,
- context1mApplied: false,
- specialSettings: [],
- cachedPriceData: undefined,
- cachedBillingModelSource: undefined,
- endpointPolicy: resolveEndpointPolicy("/v1/messages"),
- isHeaderModified: () => false,
- });
- return session as ProxySession;
- }
- function createStreamingResponse(params: {
- label: string;
- firstChunkDelayMs: number;
- controller: AbortController;
- }): Response {
- const encoder = new TextEncoder();
- let timeoutId: ReturnType<typeof setTimeout> | null = null;
- const stream = new ReadableStream<Uint8Array>({
- start(controller) {
- const onAbort = () => {
- if (timeoutId) {
- clearTimeout(timeoutId);
- }
- controller.close();
- };
- if (params.controller.signal.aborted) {
- onAbort();
- return;
- }
- params.controller.signal.addEventListener("abort", onAbort, { once: true });
- timeoutId = setTimeout(() => {
- if (params.controller.signal.aborted) {
- controller.close();
- return;
- }
- controller.enqueue(encoder.encode(`data: {"provider":"${params.label}"}\n\n`));
- controller.close();
- }, params.firstChunkDelayMs);
- },
- });
- return new Response(stream, {
- status: 200,
- headers: { "content-type": "text/event-stream" },
- });
- }
- function createDelayedFailure(params: {
- delayMs: number;
- error: Error;
- controller: AbortController;
- }): Promise<Response> {
- return new Promise((_, reject) => {
- let timeoutId: ReturnType<typeof setTimeout> | null = null;
- const rejectWithError = () => {
- if (timeoutId) {
- clearTimeout(timeoutId);
- }
- reject(params.error);
- };
- if (params.controller.signal.aborted) {
- rejectWithError();
- return;
- }
- params.controller.signal.addEventListener("abort", rejectWithError, { once: true });
- timeoutId = setTimeout(() => {
- params.controller.signal.removeEventListener("abort", rejectWithError);
- reject(params.error);
- }, params.delayMs);
- });
- }
- describe("ProxyForwarder - first-byte hedge scheduling", () => {
- beforeEach(() => {
- vi.clearAllMocks();
- });
- test("first provider exceeds first-byte threshold, second provider starts and wins by first chunk", async () => {
- vi.useFakeTimers();
- try {
- const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 });
- const provider2 = createProvider({ id: 2, name: "p2", firstByteTimeoutStreamingMs: 100 });
- const session = createSession();
- session.setProvider(provider1);
- mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(provider2);
- const doForward = vi.spyOn(
- ProxyForwarder as unknown as {
- doForward: (...args: unknown[]) => Promise<Response>;
- },
- "doForward"
- );
- const controller1 = new AbortController();
- const controller2 = new AbortController();
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller1;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p1",
- firstChunkDelayMs: 220,
- controller: controller1,
- });
- });
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller2;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p2",
- firstChunkDelayMs: 40,
- controller: controller2,
- });
- });
- const responsePromise = ProxyForwarder.send(session);
- await vi.advanceTimersByTimeAsync(100);
- expect(doForward).toHaveBeenCalledTimes(2);
- await vi.advanceTimersByTimeAsync(50);
- const response = await responsePromise;
- expect(await response.text()).toContain('"provider":"p2"');
- expect(controller1.signal.aborted).toBe(true);
- expect(controller2.signal.aborted).toBe(false);
- expect(mocks.recordFailure).not.toHaveBeenCalled();
- expect(mocks.recordSuccess).not.toHaveBeenCalled();
- expect(session.provider?.id).toBe(2);
- expect(mocks.updateSessionBindingSmart).toHaveBeenCalledWith("sess-hedge", 2, 0, false, true);
- } finally {
- vi.useRealTimers();
- }
- });
- test("characterization: hedge still launches alternative provider when maxRetryAttempts > 1", async () => {
- vi.useFakeTimers();
- try {
- const provider1 = createProvider({
- id: 1,
- name: "p1",
- maxRetryAttempts: 3,
- firstByteTimeoutStreamingMs: 100,
- });
- const provider2 = createProvider({
- id: 2,
- name: "p2",
- maxRetryAttempts: 3,
- firstByteTimeoutStreamingMs: 100,
- });
- const session = createSession();
- session.setProvider(provider1);
- mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(provider2);
- const doForward = vi.spyOn(
- ProxyForwarder as unknown as {
- doForward: (...args: unknown[]) => Promise<Response>;
- },
- "doForward"
- );
- const controller1 = new AbortController();
- const controller2 = new AbortController();
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller1;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p1",
- firstChunkDelayMs: 220,
- controller: controller1,
- });
- });
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller2;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p2",
- firstChunkDelayMs: 40,
- controller: controller2,
- });
- });
- const responsePromise = ProxyForwarder.send(session);
- await vi.advanceTimersByTimeAsync(100);
- expect(doForward).toHaveBeenCalledTimes(2);
- expect(mocks.pickRandomProviderWithExclusion).toHaveBeenCalledTimes(1);
- const chainBeforeWinner = session.getProviderChain();
- expect(chainBeforeWinner).toEqual(
- expect.arrayContaining([
- expect.objectContaining({ reason: "hedge_triggered", id: 1 }),
- expect.objectContaining({ reason: "hedge_launched", id: 2 }),
- ])
- );
- await vi.advanceTimersByTimeAsync(50);
- const response = await responsePromise;
- expect(await response.text()).toContain('"provider":"p2"');
- expect(controller1.signal.aborted).toBe(true);
- expect(session.provider?.id).toBe(2);
- } finally {
- vi.useRealTimers();
- }
- });
- test("first provider can still win after hedge started if it emits first chunk earlier than fallback", async () => {
- vi.useFakeTimers();
- try {
- const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 });
- const provider2 = createProvider({ id: 2, name: "p2", firstByteTimeoutStreamingMs: 100 });
- const session = createSession();
- session.setProvider(provider1);
- mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(provider2);
- const doForward = vi.spyOn(
- ProxyForwarder as unknown as {
- doForward: (...args: unknown[]) => Promise<Response>;
- },
- "doForward"
- );
- const controller1 = new AbortController();
- const controller2 = new AbortController();
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller1;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p1",
- firstChunkDelayMs: 140,
- controller: controller1,
- });
- });
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller2;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p2",
- firstChunkDelayMs: 120,
- controller: controller2,
- });
- });
- const responsePromise = ProxyForwarder.send(session);
- await vi.advanceTimersByTimeAsync(100);
- expect(doForward).toHaveBeenCalledTimes(2);
- await vi.advanceTimersByTimeAsync(45);
- const response = await responsePromise;
- expect(await response.text()).toContain('"provider":"p1"');
- expect(controller1.signal.aborted).toBe(false);
- expect(controller2.signal.aborted).toBe(true);
- expect(mocks.recordFailure).not.toHaveBeenCalled();
- expect(mocks.recordSuccess).not.toHaveBeenCalled();
- expect(session.provider?.id).toBe(1);
- } finally {
- vi.useRealTimers();
- }
- });
- test("when multiple providers all exceed threshold, hedge scheduler keeps expanding until a later provider wins", async () => {
- vi.useFakeTimers();
- try {
- const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 });
- const provider2 = createProvider({ id: 2, name: "p2", firstByteTimeoutStreamingMs: 100 });
- const provider3 = createProvider({ id: 3, name: "p3", firstByteTimeoutStreamingMs: 100 });
- const session = createSession();
- session.setProvider(provider1);
- mocks.pickRandomProviderWithExclusion
- .mockResolvedValueOnce(provider2)
- .mockResolvedValueOnce(provider3);
- const doForward = vi.spyOn(
- ProxyForwarder as unknown as {
- doForward: (...args: unknown[]) => Promise<Response>;
- },
- "doForward"
- );
- const controller1 = new AbortController();
- const controller2 = new AbortController();
- const controller3 = new AbortController();
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller1;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p1",
- firstChunkDelayMs: 400,
- controller: controller1,
- });
- });
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller2;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p2",
- firstChunkDelayMs: 400,
- controller: controller2,
- });
- });
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller3;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p3",
- firstChunkDelayMs: 20,
- controller: controller3,
- });
- });
- const responsePromise = ProxyForwarder.send(session);
- await vi.advanceTimersByTimeAsync(200);
- expect(doForward).toHaveBeenCalledTimes(3);
- await vi.advanceTimersByTimeAsync(25);
- const response = await responsePromise;
- expect(await response.text()).toContain('"provider":"p3"');
- expect(controller1.signal.aborted).toBe(true);
- expect(controller2.signal.aborted).toBe(true);
- expect(controller3.signal.aborted).toBe(false);
- expect(mocks.recordFailure).not.toHaveBeenCalled();
- expect(mocks.recordSuccess).not.toHaveBeenCalled();
- expect(session.provider?.id).toBe(3);
- } finally {
- vi.useRealTimers();
- }
- });
- test("client abort before any winner should abort all in-flight attempts, return 499, and clear sticky provider binding", async () => {
- vi.useFakeTimers();
- try {
- const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 });
- const provider2 = createProvider({ id: 2, name: "p2", firstByteTimeoutStreamingMs: 100 });
- const clientAbortController = new AbortController();
- const session = createSession(clientAbortController.signal);
- session.setProvider(provider1);
- mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(provider2);
- const doForward = vi.spyOn(
- ProxyForwarder as unknown as {
- doForward: (...args: unknown[]) => Promise<Response>;
- },
- "doForward"
- );
- const controller1 = new AbortController();
- const controller2 = new AbortController();
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller1;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p1",
- firstChunkDelayMs: 500,
- controller: controller1,
- });
- });
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller2;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p2",
- firstChunkDelayMs: 500,
- controller: controller2,
- });
- });
- const responsePromise = ProxyForwarder.send(session);
- const rejection = expect(responsePromise).rejects.toMatchObject({
- statusCode: 499,
- });
- await vi.advanceTimersByTimeAsync(100);
- expect(doForward).toHaveBeenCalledTimes(2);
- clientAbortController.abort(new Error("client_cancelled"));
- await vi.runAllTimersAsync();
- await rejection;
- expect(controller1.signal.aborted).toBe(true);
- expect(controller2.signal.aborted).toBe(true);
- expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge");
- expect(mocks.recordFailure).not.toHaveBeenCalled();
- expect(mocks.recordSuccess).not.toHaveBeenCalled();
- } finally {
- vi.useRealTimers();
- }
- });
- test("hedge launcher rejection should settle request instead of hanging", async () => {
- vi.useFakeTimers();
- try {
- const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 });
- const session = createSession();
- session.setProvider(provider1);
- mocks.pickRandomProviderWithExclusion.mockRejectedValueOnce(new Error("selector down"));
- const doForward = vi.spyOn(
- ProxyForwarder as unknown as {
- doForward: (...args: unknown[]) => Promise<Response>;
- },
- "doForward"
- );
- const controller1 = new AbortController();
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller1;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p1",
- firstChunkDelayMs: 500,
- controller: controller1,
- });
- });
- const responsePromise = ProxyForwarder.send(session);
- const rejection = expect(responsePromise).rejects.toMatchObject({
- statusCode: 503,
- });
- await vi.advanceTimersByTimeAsync(100);
- await vi.runAllTimersAsync();
- await rejection;
- expect(controller1.signal.aborted).toBe(true);
- } finally {
- vi.useRealTimers();
- }
- });
- test("strict endpoint pool exhaustion should converge to terminal fallback instead of provider-specific error", async () => {
- vi.useFakeTimers();
- try {
- const provider1 = createProvider({
- id: 1,
- name: "p1",
- providerType: "claude",
- providerVendorId: 123,
- firstByteTimeoutStreamingMs: 100,
- });
- const session = createSession();
- session.requestUrl = new URL("https://example.com/v1/messages");
- session.setProvider(provider1);
- mocks.getPreferredProviderEndpoints.mockRejectedValueOnce(new Error("Redis connection lost"));
- mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(null);
- const responsePromise = ProxyForwarder.send(session);
- const errorPromise = responsePromise.catch((rejection) => rejection as UpstreamProxyError);
- await vi.runAllTimersAsync();
- const error = await errorPromise;
- expect(mocks.pickRandomProviderWithExclusion).toHaveBeenCalled();
- expect(error).toBeInstanceOf(UpstreamProxyError);
- expect(error.statusCode).toBe(503);
- expect(error.message).toBe("所有供应商暂时不可用,请稍后重试");
- } finally {
- vi.useRealTimers();
- }
- });
- test.each([
- {
- name: "provider error",
- category: ProxyErrorCategory.PROVIDER_ERROR,
- errorFactory: (provider: Provider) =>
- new UpstreamProxyError("Provider returned 401: invalid key", 401, {
- body: '{"error":"invalid_api_key"}',
- providerId: provider.id,
- providerName: provider.name,
- }),
- },
- {
- name: "resource not found",
- category: ProxyErrorCategory.RESOURCE_NOT_FOUND,
- errorFactory: (provider: Provider) =>
- new UpstreamProxyError("Provider returned 404: model not found", 404, {
- body: '{"error":"model_not_found"}',
- providerId: provider.id,
- providerName: provider.name,
- }),
- },
- {
- name: "system error",
- category: ProxyErrorCategory.SYSTEM_ERROR,
- errorFactory: () => new Error("fetch failed"),
- },
- ])("when a real hedge race ends with only $name, terminal error should be generic fallback", async ({
- category,
- errorFactory,
- }) => {
- vi.useFakeTimers();
- try {
- const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 });
- const provider2 = createProvider({ id: 2, name: "p2", firstByteTimeoutStreamingMs: 100 });
- const session = createSession();
- session.setProvider(provider1);
- mocks.pickRandomProviderWithExclusion
- .mockResolvedValueOnce(provider2)
- .mockResolvedValueOnce(null);
- mocks.categorizeErrorAsync.mockResolvedValueOnce(category).mockResolvedValueOnce(category);
- const doForward = vi.spyOn(
- ProxyForwarder as unknown as {
- doForward: (...args: unknown[]) => Promise<Response>;
- },
- "doForward"
- );
- const controller1 = new AbortController();
- const controller2 = new AbortController();
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller1;
- runtime.clearResponseTimeout = vi.fn();
- return createDelayedFailure({
- delayMs: 150,
- error: errorFactory(provider1),
- controller: controller1,
- });
- });
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller2;
- runtime.clearResponseTimeout = vi.fn();
- return createDelayedFailure({
- delayMs: 160,
- error: errorFactory(provider2),
- controller: controller2,
- });
- });
- const responsePromise = ProxyForwarder.send(session);
- const errorPromise = responsePromise.catch((rejection) => rejection as UpstreamProxyError);
- await vi.advanceTimersByTimeAsync(100);
- expect(doForward).toHaveBeenCalledTimes(2);
- await vi.runAllTimersAsync();
- const error = await errorPromise;
- expect(error).toBeInstanceOf(UpstreamProxyError);
- expect(error.statusCode).toBe(503);
- expect(error.message).toBe("所有供应商暂时不可用,请稍后重试");
- expect(error.message).not.toContain("invalid key");
- expect(error.message).not.toContain("model not found");
- expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge");
- } finally {
- vi.useRealTimers();
- }
- });
- test("non-retryable client errors should stop hedge immediately and preserve original error", async () => {
- const provider1 = createProvider({ id: 1, name: "p1", firstByteTimeoutStreamingMs: 100 });
- const provider2 = createProvider({ id: 2, name: "p2", firstByteTimeoutStreamingMs: 100 });
- const session = createSession();
- session.setProvider(provider1);
- const originalError = new UpstreamProxyError("prompt too long", 400, {
- body: '{"error":"prompt_too_long"}',
- providerId: provider1.id,
- providerName: provider1.name,
- });
- mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(provider2);
- mocks.categorizeErrorAsync.mockResolvedValueOnce(ProxyErrorCategory.NON_RETRYABLE_CLIENT_ERROR);
- const doForward = vi.spyOn(
- ProxyForwarder as unknown as {
- doForward: (...args: unknown[]) => Promise<Response>;
- },
- "doForward"
- );
- doForward.mockRejectedValueOnce(originalError);
- const error = await ProxyForwarder.send(session).catch(
- (rejection) => rejection as UpstreamProxyError
- );
- expect(error).toBe(originalError);
- expect(error.message).toBe("prompt too long");
- expect(doForward).toHaveBeenCalledTimes(1);
- expect(mocks.pickRandomProviderWithExclusion).not.toHaveBeenCalled();
- expect(mocks.clearSessionProvider).toHaveBeenCalledWith("sess-hedge");
- expect(session.getProviderChain()).toEqual(
- expect.arrayContaining([
- expect.objectContaining({
- reason: "client_error_non_retryable",
- statusCode: 400,
- }),
- ])
- );
- });
- test("endpoint resolution failure should not inflate launchedProviderCount, winner gets request_success not hedge_winner", async () => {
- vi.useFakeTimers();
- try {
- const provider1 = createProvider({
- id: 1,
- name: "p1",
- providerVendorId: 123,
- firstByteTimeoutStreamingMs: 100,
- });
- const provider2 = createProvider({
- id: 2,
- name: "p2",
- providerVendorId: null,
- firstByteTimeoutStreamingMs: 100,
- });
- const session = createSession();
- session.requestUrl = new URL("https://example.com/v1/messages");
- session.setProvider(provider1);
- // Provider 1's strict endpoint resolution will fail
- mocks.getPreferredProviderEndpoints.mockRejectedValueOnce(
- new Error("Endpoint resolution failed")
- );
- // After provider 1 fails, pick provider 2 as alternative
- mocks.pickRandomProviderWithExclusion.mockResolvedValueOnce(provider2);
- const doForward = vi.spyOn(
- ProxyForwarder as unknown as {
- doForward: (...args: unknown[]) => Promise<Response>;
- },
- "doForward"
- );
- const controller2 = new AbortController();
- // Only provider 2 reaches doForward (provider 1 fails at endpoint resolution)
- doForward.mockImplementationOnce(async (attemptSession) => {
- const runtime = attemptSession as ProxySession & AttemptRuntime;
- runtime.responseController = controller2;
- runtime.clearResponseTimeout = vi.fn();
- return createStreamingResponse({
- label: "p2",
- firstChunkDelayMs: 10,
- controller: controller2,
- });
- });
- const responsePromise = ProxyForwarder.send(session);
- await vi.advanceTimersByTimeAsync(200);
- const response = await responsePromise;
- expect(await response.text()).toContain('"provider":"p2"');
- expect(session.provider?.id).toBe(2);
- // Key assertion: since only provider 2 actually launched (provider 1 failed at
- // endpoint resolution before incrementing launchedProviderCount), the winner
- // should be classified as "request_success" not "hedge_winner".
- const chain = session.getProviderChain();
- const winnerEntry = chain.find(
- (entry) => entry.reason === "request_success" || entry.reason === "hedge_winner"
- );
- expect(winnerEntry).toBeDefined();
- expect(winnerEntry!.reason).toBe("request_success");
- } finally {
- vi.useRealTimers();
- }
- });
- });
|