File size: 4,439 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 | import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import { initDb } from '../../db/index.js';
import { runFallbackLoop, newFallbackState } from '../../lib/fallback-loop.js';
import { acquireLease, releaseLease, resetLeases, inFlightForKey } from '../../services/ratelimit.js';
import type { RouteResult } from '../../services/router.js';
/**
* The loop must free a route's in-flight lease on every exit path. A leak would
* leave the key's concurrency budget permanently short (until the age-out
* backstop), so this covers each way an iteration can end.
*/
function leasedRoute(keyId: number): RouteResult {
const id = acquireLease('groq', 'llama-3.3-70b', keyId, 100);
return {
provider: {} as any,
modelId: 'llama-3.3-70b',
modelDbId: 1,
apiKey: 'test-key',
keyId,
platform: 'groq',
displayName: 'Llama 3.3 70B',
rpdLimit: null,
tpdLimit: null,
release: () => releaseLease(id),
};
}
function hooks(overrides: Record<string, unknown> = {}) {
return {
state: newFallbackState(),
attemptLog: [],
route: () => leasedRoute(1),
dispatch: async () => 'done' as const,
logFailure: () => undefined,
onFatal: () => undefined,
onRoutingExhausted: () => undefined,
onExhausted: () => undefined,
...overrides,
} as any;
}
describe('fallback loop releases in-flight leases', () => {
beforeEach(() => {
initDb(':memory:');
resetLeases();
vi.spyOn(console, 'log').mockImplementation(() => {});
vi.spyOn(console, 'error').mockImplementation(() => {});
vi.spyOn(console, 'warn').mockImplementation(() => {});
});
afterEach(() => {
resetLeases();
vi.restoreAllMocks();
});
it('releases after a successful dispatch', async () => {
await runFallbackLoop(hooks());
expect(inFlightForKey('groq', 1)).toBe(0);
});
it('releases after a committed stream', async () => {
await runFallbackLoop(hooks({ dispatch: async () => 'committed' as const }));
expect(inFlightForKey('groq', 1)).toBe(0);
});
it('releases every attempt when retryable failures exhaust the budget', async () => {
let dispatched = 0;
await runFallbackLoop(hooks({
maxRetries: 4,
dispatch: async () => {
dispatched++;
throw Object.assign(new Error('Groq API error 429: rate limit'), { status: 429 });
},
}));
expect(dispatched).toBe(4);
// Four leases were taken across the four attempts; none may survive.
expect(inFlightForKey('groq', 1)).toBe(0);
});
it('releases on an auth failure that rotates the key', async () => {
await runFallbackLoop(hooks({
maxRetries: 2,
dispatch: async () => {
throw Object.assign(new Error('Groq API error 401: Invalid API Key'), { status: 401 });
},
}));
expect(inFlightForKey('groq', 1)).toBe(0);
});
it('releases on a non-retryable fatal error', async () => {
const onFatal = vi.fn();
await runFallbackLoop(hooks({
onFatal,
dispatch: async () => { throw Object.assign(new Error('boom'), { status: 418 }); },
}));
expect(onFatal).toHaveBeenCalled();
expect(inFlightForKey('groq', 1)).toBe(0);
});
it('releases when dispatch violates the outcome contract', async () => {
const onFatal = vi.fn();
await runFallbackLoop(hooks({ onFatal, dispatch: async () => 'nonsense' as any }));
expect(onFatal).toHaveBeenCalled();
expect(inFlightForKey('groq', 1)).toBe(0);
});
it('does not mask the provider error if release itself throws', async () => {
const onFatal = vi.fn();
const bad = {
...leasedRoute(1),
release: () => { throw new Error('release exploded'); },
};
await runFallbackLoop(hooks({
route: () => bad,
onFatal,
dispatch: async () => { throw Object.assign(new Error('real provider failure'), { status: 418 }); },
})).catch(() => undefined);
// onFatal must have seen the provider error, not the release error.
expect(onFatal).toHaveBeenCalled();
const err = onFatal.mock.calls[0]![1] as Error;
expect(err.message).toContain('real provider failure');
});
it('tolerates a route with no release function', async () => {
const noRelease = { ...leasedRoute(2) } as RouteResult;
delete (noRelease as { release?: unknown }).release;
await expect(runFallbackLoop(hooks({ route: () => noRelease }))).resolves.toBeUndefined();
});
});
|