File size: 5,937 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 | import { describe, it, expect, beforeAll, beforeEach, vi } from 'vitest';
// Model-level failure benching: a model that keeps failing upstream must sink
// out of routing on EVERY key that can serve it, not only on the key that
// happened to fail last — the failure window counts across keys, so a per-key
// bench would leave the siblings feeding the same sick model.
vi.mock('../../services/health.js', () => ({
checkKeyHealth: vi.fn(),
// recordUpstreamSuccess touches this on the streak-clearing path.
markKeyHealthyFromRequest: vi.fn(),
}));
import { initDb, getDb } from '../../db/index.js';
import { encrypt } from '../../lib/crypto.js';
import {
newFallbackState,
recordRetryableFailure,
recordUpstreamSuccess,
MODEL_FAILURE_COOLDOWN_MS,
MODEL_FAILURE_THRESHOLD,
MODEL_FAILURE_WINDOW_MS,
resetModelFailureWindows,
} from '../../lib/fallback-loop.js';
import {
getActiveCooldownsForKeys,
isOnCooldown,
resetKeyLocalityCache,
} from '../../services/ratelimit.js';
import type { RouteResult } from '../../services/router.js';
const PLATFORM = 'groq';
let keyA = 0;
let keyB = 0;
let modelA: { id: number; model_id: string };
let modelB: { id: number; model_id: string };
function insertKey(label: string): number {
const { encrypted, iv, authTag } = encrypt(`test-${label}`);
const info = getDb().prepare(`
INSERT INTO api_keys (platform, label, encrypted_key, iv, auth_tag, status, enabled)
VALUES (?, ?, ?, ?, ?, 'healthy', 1)
`).run(PLATFORM, label, encrypted, iv, authTag);
return Number(info.lastInsertRowid);
}
function routeFor(model: { id: number; model_id: string }, keyId: number): RouteResult {
return {
provider: {} as any,
modelId: model.model_id,
modelDbId: model.id,
apiKey: 'k',
keyId,
platform: PLATFORM,
displayName: model.model_id,
rpdLimit: null,
tpdLimit: null,
};
}
// A plain 5xx: retryable, no quota signal, so the per-key bench is the short
// transient one (90s) — far below the model bench, which is what makes the
// assertions below able to tell the two apart.
const upstream500 = () => Object.assign(new Error('Groq API error 500: upstream'), { status: 500 });
function failOnce(model: { id: number; model_id: string }, keyId: number, now?: number) {
// Fresh state per call: each failure stands for a separate request.
recordRetryableFailure(routeFor(model, keyId), upstream500(), newFallbackState(), now);
}
function cooldownFor(model: { id: number; model_id: string }, keyId: number) {
return getActiveCooldownsForKeys([keyId])
.get(keyId)
?.find(c => c.platform === PLATFORM && c.modelId === model.model_id);
}
beforeAll(() => {
process.env.ENCRYPTION_KEY = '0'.repeat(64);
initDb(':memory:');
const db = getDb();
db.prepare('DELETE FROM api_keys').run();
keyA = insertKey('key-a');
keyB = insertKey('key-b');
resetKeyLocalityCache();
const models = db.prepare(
'SELECT id, model_id FROM models WHERE platform = ? ORDER BY id LIMIT 2'
).all(PLATFORM) as { id: number; model_id: string }[];
expect(models.length).toBe(2);
[modelA, modelB] = models;
});
beforeEach(() => {
getDb().prepare('DELETE FROM rate_limit_cooldowns').run();
// The failure windows are module state and outlive the DB wipe above, so a
// case would otherwise start partway to the threshold.
resetModelFailureWindows();
});
describe('model-level failure benching covers every key of the model', () => {
it('benches BOTH keys once the window trips, even though only one key failed last', () => {
// Three failures inside the window, spread across the two keys.
failOnce(modelA, keyA);
failOnce(modelA, keyB);
expect(MODEL_FAILURE_THRESHOLD).toBe(3);
failOnce(modelA, keyA);
for (const keyId of [keyA, keyB]) {
expect(isOnCooldown(PLATFORM, modelA.model_id, keyId)).toBe(true);
const cd = cooldownFor(modelA, keyId);
expect(cd).toBeDefined();
// The per-key transient bench is 90s; only the model bench reaches here.
expect(cd!.remainingMs).toBeGreaterThan(MODEL_FAILURE_COOLDOWN_MS / 2);
}
});
it('starts the streak over after a served request, so the next failure benches nothing extra', () => {
failOnce(modelB, keyA);
failOnce(modelB, keyB);
// A served request is counter-evidence: the two failures above no longer count.
recordUpstreamSuccess(routeFor(modelB, keyA), 100);
failOnce(modelB, keyA);
// Only the short per-key benches, no model-wide one.
const benchedB = cooldownFor(modelB, keyB);
expect(benchedB).toBeDefined();
expect(benchedB!.remainingMs).toBeLessThan(MODEL_FAILURE_COOLDOWN_MS / 2);
const benchedA = cooldownFor(modelB, keyA);
expect(benchedA).toBeDefined();
expect(benchedA!.remainingMs).toBeLessThan(MODEL_FAILURE_COOLDOWN_MS / 2);
});
// The window is a SLIDING one, which is only observable if the clock can be
// moved. Without an injectable `now` this case would have to sleep for 15
// real minutes, so the behaviour was never pinned.
it('does not trip when the failures are spread beyond the sliding window', () => {
const t0 = Date.now();
failOnce(modelA, keyA, t0);
failOnce(modelA, keyA, t0 + MODEL_FAILURE_WINDOW_MS + 1);
failOnce(modelA, keyA, t0 + 2 * (MODEL_FAILURE_WINDOW_MS + 1));
// Three failures, but never three inside one window: no model-wide bench.
const benched = cooldownFor(modelA, keyB);
expect(benched).toBeUndefined();
});
it('trips when the same three failures fall inside one window', () => {
const t0 = Date.now();
failOnce(modelA, keyA, t0);
failOnce(modelA, keyA, t0 + 1000);
failOnce(modelA, keyA, t0 + 2000);
// keyB never failed, but the model-wide bench covers it.
const benched = cooldownFor(modelA, keyB);
expect(benched).toBeDefined();
expect(benched!.remainingMs).toBeGreaterThan(MODEL_FAILURE_COOLDOWN_MS / 2);
});
});
|