File size: 11,749 Bytes
dcedb25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 | 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);
});
});
|