Spaces:

automindy
/

fokemal

Sleeping

fokemal / server /src /__tests__ /lib /fallback-loop-hedge.test.ts

Sync to upstream v0.13.7, keep HF Space adaptations

dcedb25 22 days ago

11.7 kB

import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import { initDb, getDb } from '../../db/index.js';
import { runFallbackLoop, newFallbackState, type FallbackState } from '../../lib/fallback-loop.js';
import { isHedgeAbortError, newHedgeAbortError, isClientAbortError, isRetryableError } from '../../lib/error-classify.js';
import { acquireLease, releaseLease, resetLeases } from '../../services/ratelimit.js';
import type { RouteResult } from '../../services/router.js';

/**
 * Fallback-v2 hedging: when the wall-clock retry budget expires MID-FLIGHT,
 * the loop calls abortInFlight() so a stalled attempt is abandoned instead of
 * only refusing to start the next retry behind it. The surface aborts its
 * composed fetch signal with newHedgeAbortError(); the loop must render
 * timedOut exhaustion, and the provider-health bookkeeping — cooldown, skip
 * entry, logFailure row — depends on WHOSE silence spent the budget: an attempt
 * handed only the scraps a ladder left is not benched, one that stayed silent
 * for most of the whole budget is.
 */

let keySeq = 900;
function leasedRoute(): { route: RouteResult; keyId: number } {
  const keyId = ++keySeq;
  const id = acquireLease('fake', 'fake-model', keyId, 100);
  return {
    keyId,
    route: {
      provider: {} as any,
      modelId: 'fake-model',
      modelDbId: 900000 + keyId,
      apiKey: 'test-key',
      keyId,
      platform: 'fake',
      displayName: 'Fake Model',
      rpdLimit: null,
      tpdLimit: null,
      release: () => releaseLease(id),
    },
  };
}

function hooks(state: FallbackState, overrides: Record<string, unknown>) {
  return {
    state,
    timeBudgetMs: 0,
    logFailure: vi.fn(),
    onFatal: vi.fn(),
    onRoutingExhausted: vi.fn(),
    onExhausted: vi.fn(),
    ...overrides,
  } as any;
}

// An upstream that hangs until the composed signal aborts, then rejects with
// the marked hedge error — mirrors what a real fetch does when the surface
// aborts its signal via abortInFlight().
function stalledDispatch(hedgeAbort: AbortController) {
  return async (): Promise<never> => {
    return await new Promise((_resolve, reject) => {
      hedgeAbort.signal.addEventListener('abort', () => reject(newHedgeAbortError()), { once: true });
    });
  };
}

describe('fallback loop time-budget hedging', () => {
  beforeEach(() => {
    initDb(':memory:');
    resetLeases();
    vi.spyOn(console, 'log').mockImplementation(() => {});
  });

  afterEach(() => {
    resetLeases();
    vi.restoreAllMocks();
  });

  it('classifies the hedge error as budget-abort, not retryable or client-abort', () => {
    const err = newHedgeAbortError();
    expect(isHedgeAbortError(err)).toBe(true);
    // Must never enter the retryable ladder as a provider failure.
    expect(isRetryableError(err)).toBe(false);
    // Distinct from the client-disconnect abort — different branch in the loop.
    expect(isClientAbortError(err)).toBe(false);
  });

  it('aborts the stalled in-flight attempt when the budget expires, rendering timedOut exhaustion', async () => {
    const state = newFallbackState();
    const hedgeAbort = new AbortController();
    // First two attempts fail quickly and retryably (each ~30ms < the 100ms
    // budget, so the third attempt gets to start); the third hangs and is
    // aborted when the remaining budget (~40ms) expires mid-flight.
    const dispatch = vi.fn()
      .mockImplementationOnce(async () => {
        await new Promise(r => setTimeout(r, 30));
        throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
      })
      .mockImplementationOnce(async () => {
        await new Promise(r => setTimeout(r, 30));
        throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
      })
      .mockImplementation(stalledDispatch(hedgeAbort));

    const h = hooks(state, {
      maxRetries: 5,
      timeBudgetMs: 100,
      abortInFlight: () => hedgeAbort.abort(newHedgeAbortError()),
      route: () => leasedRoute().route,
      dispatch,
    });

    await runFallbackLoop(h);

    // The stalled attempt was abandoned mid-flight; the budget is spent, so
    // no fourth attempt starts.
    expect(dispatch).toHaveBeenCalledTimes(3);
    expect(h.onExhausted).toHaveBeenCalledTimes(1);
    expect(h.onExhausted.mock.calls[0][1].timedOut).toBe(true);
    // The two quick 429 failures were logged as failures; the hedge abort of
    // the third attempt must NOT add a third logFailure row, and there is no
    // fatal render or routing-exhaustion render.
    expect(h.logFailure).toHaveBeenCalledTimes(2);
    expect(h.onFatal).not.toHaveBeenCalled();
    expect(h.onRoutingExhausted).not.toHaveBeenCalled();
  });

  it('the hedge abort adds no cooldown, no skip, no extra failure beyond the retryable attempts', async () => {
    const state = newFallbackState();
    const hedgeAbort = new AbortController();
    // Same ladder shape as above: two quick retryable failures let the third
    // attempt start, then the budget expires mid-flight and the hedge abort
    // abandons it. Hedging only arms for attempt > 1 (the first attempt and
    // the first retry always run — #751).
    const dispatch = vi.fn()
      .mockImplementationOnce(async () => {
        await new Promise(r => setTimeout(r, 30));
        throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
      })
      .mockImplementationOnce(async () => {
        await new Promise(r => setTimeout(r, 30));
        throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
      })
      .mockImplementation(stalledDispatch(hedgeAbort));

    const h = hooks(state, {
      maxRetries: 5,
      timeBudgetMs: 100,
      abortInFlight: () => hedgeAbort.abort(newHedgeAbortError()),
      route: () => leasedRoute().route,
      dispatch,
    });
    await runFallbackLoop(h);

    // The two 429s each produced their normal failure bookkeeping (a cooldown
    // row and a skip entry); the hedge abort must NOT add a third of either —
    // the budget expiring is not a provider-health signal.
    const cooldowns = getDb().prepare('SELECT COUNT(*) AS n FROM rate_limit_cooldowns').get() as { n: number };
    expect(cooldowns.n).toBe(2);
    expect(state.skipKeys.size).toBe(2);
    expect(state.skipModels.size).toBe(0);
    expect(h.logFailure).toHaveBeenCalledTimes(2);
    // And the request rendered as timedOut exhaustion, not a third failure.
    expect(h.onExhausted).toHaveBeenCalledTimes(1);
    expect(h.onExhausted.mock.calls[0][1].timedOut).toBe(true);
  });

  it('benches the route when one attempt stayed silent for the whole budget', async () => {
    // The production shape behind this branch. Same ladder position as the two
    // tests above — hedging only arms from the third attempt on (#751) — but
    // here the two failures ahead of it are INSTANT, so the stalled attempt
    // inherits essentially the whole budget and then sends no first byte for
    // all of it. That is a stalled upstream, not a scheduling accident:
    // unbenched it stays at the head of the route order and re-stalls every
    // subsequent request, burning the full budget each time while the healthy
    // routes queued behind it never run.
    const state = newFallbackState();
    const hedgeAbort = new AbortController();
    const fast429 = async () => {
      throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
    };
    const dispatch = vi.fn()
      .mockImplementationOnce(fast429)
      .mockImplementationOnce(fast429)
      .mockImplementation(stalledDispatch(hedgeAbort));

    const h = hooks(state, {
      maxRetries: 5,
      timeBudgetMs: 100,
      abortInFlight: () => hedgeAbort.abort(newHedgeAbortError()),
      route: () => leasedRoute().route,
      dispatch,
    });

    await runFallbackLoop(h);

    // Still a timedOut exhaustion — the render is unchanged.
    expect(h.onExhausted).toHaveBeenCalledTimes(1);
    expect(h.onExhausted.mock.calls[0][1].timedOut).toBe(true);
    // But the stalled route now carries a cooldown and a skip entry of its own,
    // on top of the two fast 429s', so the next request routes past it. This is
    // the exact assertion that inverts for the shared-budget case above, where
    // the same three attempts produce only two of each.
    const cooldowns = getDb().prepare('SELECT COUNT(*) AS n FROM rate_limit_cooldowns').get() as { n: number };
    expect(cooldowns.n).toBe(3);
    expect(state.skipKeys.size).toBe(3);
    expect(h.logFailure).toHaveBeenCalledTimes(3);
  });

  it('never cancels an attempt that already committed, however long it then runs', async () => {
    // The whole point of the hedge is to kill a SILENT attempt. Once a stream
    // has flushed its headers the answer is already reaching the client and the
    // loop could not fail over anyway, so cancelling would truncate a healthy
    // response for nothing. Slow != stalled.
    const state = newFallbackState();
    const abortInFlight = vi.fn();
    const dispatch = vi.fn()
      .mockImplementationOnce(async () => {
        await new Promise(r => setTimeout(r, 30));
        throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
      })
      .mockImplementationOnce(async () => {
        await new Promise(r => setTimeout(r, 30));
        throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
      })
      // Third attempt commits immediately, then keeps streaming well past the
      // 100ms budget — exactly the slow-but-healthy generation that used to be
      // truncated at the budget.
      .mockImplementation(async (_route: unknown, _attempt: number, ctx: { disarmHedge: () => void }) => {
        ctx.disarmHedge();
        await new Promise(r => setTimeout(r, 250));
        return 'committed';
      });

    const h = hooks(state, {
      maxRetries: 5,
      timeBudgetMs: 100,
      abortInFlight,
      route: () => leasedRoute().route,
      dispatch,
    });

    await runFallbackLoop(h);

    expect(dispatch).toHaveBeenCalledTimes(3);
    // The committed attempt outlived the budget and was left alone.
    expect(abortInFlight).not.toHaveBeenCalled();
    expect(h.onExhausted).not.toHaveBeenCalled();
    expect(h.onFatal).not.toHaveBeenCalled();
    // Only the two real 429s were booked as failures.
    expect(h.logFailure).toHaveBeenCalledTimes(2);
  });

  it('still aborts when the attempt stays silent past the budget', async () => {
    // The complement of the test above: no disarmHedge call (nothing was ever
    // flushed), so the stalled attempt is still cancelled.
    const state = newFallbackState();
    const hedgeAbort = new AbortController();
    const abortInFlight = vi.fn(() => hedgeAbort.abort(newHedgeAbortError()));
    const dispatch = vi.fn()
      .mockImplementationOnce(async () => {
        await new Promise(r => setTimeout(r, 30));
        throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
      })
      .mockImplementationOnce(async () => {
        await new Promise(r => setTimeout(r, 30));
        throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
      })
      .mockImplementation(stalledDispatch(hedgeAbort));

    const h = hooks(state, {
      maxRetries: 5,
      timeBudgetMs: 100,
      abortInFlight,
      route: () => leasedRoute().route,
      dispatch,
    });

    await runFallbackLoop(h);

    expect(abortInFlight).toHaveBeenCalledTimes(1);
    expect(h.onExhausted).toHaveBeenCalledTimes(1);
    expect(h.onExhausted.mock.calls[0][1].timedOut).toBe(true);
  });
});

	import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
	import { initDb, getDb } from '../../db/index.js';
	import { runFallbackLoop, newFallbackState, type FallbackState } from '../../lib/fallback-loop.js';
	import { isHedgeAbortError, newHedgeAbortError, isClientAbortError, isRetryableError } from '../../lib/error-classify.js';
	import { acquireLease, releaseLease, resetLeases } from '../../services/ratelimit.js';
	import type { RouteResult } from '../../services/router.js';

	/**
	* Fallback-v2 hedging: when the wall-clock retry budget expires MID-FLIGHT,
	* the loop calls abortInFlight() so a stalled attempt is abandoned instead of
	* only refusing to start the next retry behind it. The surface aborts its
	* composed fetch signal with newHedgeAbortError(); the loop must render
	* timedOut exhaustion, and the provider-health bookkeeping — cooldown, skip
	* entry, logFailure row — depends on WHOSE silence spent the budget: an attempt
	* handed only the scraps a ladder left is not benched, one that stayed silent
	* for most of the whole budget is.
	*/

	let keySeq = 900;
	function leasedRoute(): { route: RouteResult; keyId: number } {
	const keyId = ++keySeq;
	const id = acquireLease('fake', 'fake-model', keyId, 100);
	return {
	keyId,
	route: {
	provider: {} as any,
	modelId: 'fake-model',
	modelDbId: 900000 + keyId,
	apiKey: 'test-key',
	keyId,
	platform: 'fake',
	displayName: 'Fake Model',
	rpdLimit: null,
	tpdLimit: null,
	release: () => releaseLease(id),
	},
	};
	}

	function hooks(state: FallbackState, overrides: Record<string, unknown>) {
	return {
	state,
	timeBudgetMs: 0,
	logFailure: vi.fn(),
	onFatal: vi.fn(),
	onRoutingExhausted: vi.fn(),
	onExhausted: vi.fn(),
	...overrides,
	} as any;
	}

	// An upstream that hangs until the composed signal aborts, then rejects with
	// the marked hedge error — mirrors what a real fetch does when the surface
	// aborts its signal via abortInFlight().
	function stalledDispatch(hedgeAbort: AbortController) {
	return async (): Promise<never> => {
	return await new Promise((_resolve, reject) => {
	hedgeAbort.signal.addEventListener('abort', () => reject(newHedgeAbortError()), { once: true });
	});
	};
	}

	describe('fallback loop time-budget hedging', () => {
	beforeEach(() => {
	initDb(':memory:');
	resetLeases();
	vi.spyOn(console, 'log').mockImplementation(() => {});
	});

	afterEach(() => {
	resetLeases();
	vi.restoreAllMocks();
	});

	it('classifies the hedge error as budget-abort, not retryable or client-abort', () => {
	const err = newHedgeAbortError();
	expect(isHedgeAbortError(err)).toBe(true);
	// Must never enter the retryable ladder as a provider failure.
	expect(isRetryableError(err)).toBe(false);
	// Distinct from the client-disconnect abort — different branch in the loop.
	expect(isClientAbortError(err)).toBe(false);
	});

	it('aborts the stalled in-flight attempt when the budget expires, rendering timedOut exhaustion', async () => {
	const state = newFallbackState();
	const hedgeAbort = new AbortController();
	// First two attempts fail quickly and retryably (each ~30ms < the 100ms
	// budget, so the third attempt gets to start); the third hangs and is
	// aborted when the remaining budget (~40ms) expires mid-flight.
	const dispatch = vi.fn()
	.mockImplementationOnce(async () => {
	await new Promise(r => setTimeout(r, 30));
	throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
	})
	.mockImplementationOnce(async () => {
	await new Promise(r => setTimeout(r, 30));
	throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
	})
	.mockImplementation(stalledDispatch(hedgeAbort));

	const h = hooks(state, {
	maxRetries: 5,
	timeBudgetMs: 100,
	abortInFlight: () => hedgeAbort.abort(newHedgeAbortError()),
	route: () => leasedRoute().route,
	dispatch,
	});

	await runFallbackLoop(h);

	// The stalled attempt was abandoned mid-flight; the budget is spent, so
	// no fourth attempt starts.
	expect(dispatch).toHaveBeenCalledTimes(3);
	expect(h.onExhausted).toHaveBeenCalledTimes(1);
	expect(h.onExhausted.mock.calls[0][1].timedOut).toBe(true);
	// The two quick 429 failures were logged as failures; the hedge abort of
	// the third attempt must NOT add a third logFailure row, and there is no
	// fatal render or routing-exhaustion render.
	expect(h.logFailure).toHaveBeenCalledTimes(2);
	expect(h.onFatal).not.toHaveBeenCalled();
	expect(h.onRoutingExhausted).not.toHaveBeenCalled();
	});

	it('the hedge abort adds no cooldown, no skip, no extra failure beyond the retryable attempts', async () => {
	const state = newFallbackState();
	const hedgeAbort = new AbortController();
	// Same ladder shape as above: two quick retryable failures let the third
	// attempt start, then the budget expires mid-flight and the hedge abort
	// abandons it. Hedging only arms for attempt > 1 (the first attempt and
	// the first retry always run — #751).
	const dispatch = vi.fn()
	.mockImplementationOnce(async () => {
	await new Promise(r => setTimeout(r, 30));
	throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
	})
	.mockImplementationOnce(async () => {
	await new Promise(r => setTimeout(r, 30));
	throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
	})
	.mockImplementation(stalledDispatch(hedgeAbort));

	const h = hooks(state, {
	maxRetries: 5,
	timeBudgetMs: 100,
	abortInFlight: () => hedgeAbort.abort(newHedgeAbortError()),
	route: () => leasedRoute().route,
	dispatch,
	});
	await runFallbackLoop(h);

	// The two 429s each produced their normal failure bookkeeping (a cooldown
	// row and a skip entry); the hedge abort must NOT add a third of either —
	// the budget expiring is not a provider-health signal.
	const cooldowns = getDb().prepare('SELECT COUNT(*) AS n FROM rate_limit_cooldowns').get() as { n: number };
	expect(cooldowns.n).toBe(2);
	expect(state.skipKeys.size).toBe(2);
	expect(state.skipModels.size).toBe(0);
	expect(h.logFailure).toHaveBeenCalledTimes(2);
	// And the request rendered as timedOut exhaustion, not a third failure.
	expect(h.onExhausted).toHaveBeenCalledTimes(1);
	expect(h.onExhausted.mock.calls[0][1].timedOut).toBe(true);
	});

	it('benches the route when one attempt stayed silent for the whole budget', async () => {
	// The production shape behind this branch. Same ladder position as the two
	// tests above — hedging only arms from the third attempt on (#751) — but
	// here the two failures ahead of it are INSTANT, so the stalled attempt
	// inherits essentially the whole budget and then sends no first byte for
	// all of it. That is a stalled upstream, not a scheduling accident:
	// unbenched it stays at the head of the route order and re-stalls every
	// subsequent request, burning the full budget each time while the healthy
	// routes queued behind it never run.
	const state = newFallbackState();
	const hedgeAbort = new AbortController();
	const fast429 = async () => {
	throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
	};
	const dispatch = vi.fn()
	.mockImplementationOnce(fast429)
	.mockImplementationOnce(fast429)
	.mockImplementation(stalledDispatch(hedgeAbort));

	const h = hooks(state, {
	maxRetries: 5,
	timeBudgetMs: 100,
	abortInFlight: () => hedgeAbort.abort(newHedgeAbortError()),
	route: () => leasedRoute().route,
	dispatch,
	});

	await runFallbackLoop(h);

	// Still a timedOut exhaustion — the render is unchanged.
	expect(h.onExhausted).toHaveBeenCalledTimes(1);
	expect(h.onExhausted.mock.calls[0][1].timedOut).toBe(true);
	// But the stalled route now carries a cooldown and a skip entry of its own,
	// on top of the two fast 429s', so the next request routes past it. This is
	// the exact assertion that inverts for the shared-budget case above, where
	// the same three attempts produce only two of each.
	const cooldowns = getDb().prepare('SELECT COUNT(*) AS n FROM rate_limit_cooldowns').get() as { n: number };
	expect(cooldowns.n).toBe(3);
	expect(state.skipKeys.size).toBe(3);
	expect(h.logFailure).toHaveBeenCalledTimes(3);
	});

	it('never cancels an attempt that already committed, however long it then runs', async () => {
	// The whole point of the hedge is to kill a SILENT attempt. Once a stream
	// has flushed its headers the answer is already reaching the client and the
	// loop could not fail over anyway, so cancelling would truncate a healthy
	// response for nothing. Slow != stalled.
	const state = newFallbackState();
	const abortInFlight = vi.fn();
	const dispatch = vi.fn()
	.mockImplementationOnce(async () => {
	await new Promise(r => setTimeout(r, 30));
	throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
	})
	.mockImplementationOnce(async () => {
	await new Promise(r => setTimeout(r, 30));
	throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
	})
	// Third attempt commits immediately, then keeps streaming well past the
	// 100ms budget — exactly the slow-but-healthy generation that used to be
	// truncated at the budget.
	.mockImplementation(async (_route: unknown, _attempt: number, ctx: { disarmHedge: () => void }) => {
	ctx.disarmHedge();
	await new Promise(r => setTimeout(r, 250));
	return 'committed';
	});

	const h = hooks(state, {
	maxRetries: 5,
	timeBudgetMs: 100,
	abortInFlight,
	route: () => leasedRoute().route,
	dispatch,
	});

	await runFallbackLoop(h);

	expect(dispatch).toHaveBeenCalledTimes(3);
	// The committed attempt outlived the budget and was left alone.
	expect(abortInFlight).not.toHaveBeenCalled();
	expect(h.onExhausted).not.toHaveBeenCalled();
	expect(h.onFatal).not.toHaveBeenCalled();
	// Only the two real 429s were booked as failures.
	expect(h.logFailure).toHaveBeenCalledTimes(2);
	});

	it('still aborts when the attempt stays silent past the budget', async () => {
	// The complement of the test above: no disarmHedge call (nothing was ever
	// flushed), so the stalled attempt is still cancelled.
	const state = newFallbackState();
	const hedgeAbort = new AbortController();
	const abortInFlight = vi.fn(() => hedgeAbort.abort(newHedgeAbortError()));
	const dispatch = vi.fn()
	.mockImplementationOnce(async () => {
	await new Promise(r => setTimeout(r, 30));
	throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
	})
	.mockImplementationOnce(async () => {
	await new Promise(r => setTimeout(r, 30));
	throw Object.assign(new Error('fake API error 429: rate limit'), { status: 429 });
	})
	.mockImplementation(stalledDispatch(hedgeAbort));

	const h = hooks(state, {
	maxRetries: 5,
	timeBudgetMs: 100,
	abortInFlight,
	route: () => leasedRoute().route,
	dispatch,
	});

	await runFallbackLoop(h);

	expect(abortInFlight).toHaveBeenCalledTimes(1);
	expect(h.onExhausted).toHaveBeenCalledTimes(1);
	expect(h.onExhausted.mock.calls[0][1].timedOut).toBe(true);
	});
	});