307 lines
12 KiB
TypeScript
307 lines
12 KiB
TypeScript
/**
|
||
* tests/unit/quota-saturation-token-headers.test.ts
|
||
*
|
||
* #2 — Proactive saturation from TOKEN rate-limit headers (universal).
|
||
*
|
||
* storeRateLimitHeaders previously only captured the per-minute REQUEST headers
|
||
* (anthropic-ratelimit-requests-* / x-ratelimit-*-requests). TOKEN headers ride
|
||
* on EVERY upstream response (success too), so they let us throttle proactively
|
||
* before a 429.
|
||
*
|
||
* This suite asserts:
|
||
* - storeRateLimitHeaders parses Anthropic token headers
|
||
* (anthropic-ratelimit-tokens-{limit,remaining,reset}, reset = RFC3339)
|
||
* and OpenAI token headers (x-ratelimit-{limit,remaining,reset}-tokens,
|
||
* reset = duration string like "6m0s"/"1s").
|
||
* - the resulting token-header saturation = 1 − remaining/limit (clamped 0..1).
|
||
* - reset is normalized to an epoch-ms `resetAt` (RFC3339 → epoch,
|
||
* duration → now + parsed seconds).
|
||
* - the existing REQUEST path is untouched (no regression).
|
||
* - getSaturation for a generic provider (openai) surfaces the token-header
|
||
* signal when usage returns nothing (fallback/complement), failing open to 0.
|
||
*
|
||
* Pure parse + cache; no DB, no network. The 30s saturation cache and the
|
||
* token-header cache are cleared between asserts.
|
||
*/
|
||
|
||
import test from "node:test";
|
||
import assert from "node:assert/strict";
|
||
|
||
const satMod = await import("../../src/lib/quota/saturationSignals.ts");
|
||
const {
|
||
storeRateLimitHeaders,
|
||
getTokenHeaderSaturation,
|
||
getSaturation,
|
||
_clearSaturationCache,
|
||
_clearRateLimitHeaders,
|
||
__setGenericUsageFetcherForTests,
|
||
} = satMod;
|
||
|
||
test.afterEach(() => {
|
||
_clearSaturationCache();
|
||
_clearRateLimitHeaders();
|
||
__setGenericUsageFetcherForTests(null);
|
||
});
|
||
|
||
// ─── Anthropic token headers (reset = RFC3339) ───────────────────────────────
|
||
|
||
test("Anthropic token headers → saturation = 1 − remaining/limit", () => {
|
||
_clearRateLimitHeaders();
|
||
// 1000 limit, 250 remaining → 750 used → 0.75 saturation.
|
||
storeRateLimitHeaders("anth-conn-1", "anthropic", {
|
||
"anthropic-ratelimit-tokens-limit": "1000",
|
||
"anthropic-ratelimit-tokens-remaining": "250",
|
||
"anthropic-ratelimit-tokens-reset": "2026-01-01T00:00:30Z",
|
||
});
|
||
|
||
const sig = getTokenHeaderSaturation("anthropic", "anth-conn-1");
|
||
assert.ok(sig, "expected a token-header saturation signal");
|
||
assert.ok(Math.abs(sig!.saturation - 0.75) < 1e-9, `expected ≈0.75, got ${sig!.saturation}`);
|
||
});
|
||
|
||
test("Anthropic token reset (RFC3339) is normalized to epoch ms", () => {
|
||
_clearRateLimitHeaders();
|
||
storeRateLimitHeaders("anth-conn-2", "anthropic", {
|
||
"anthropic-ratelimit-tokens-limit": "1000",
|
||
"anthropic-ratelimit-tokens-remaining": "0",
|
||
"anthropic-ratelimit-tokens-reset": "2026-01-01T00:00:30Z",
|
||
});
|
||
|
||
const sig = getTokenHeaderSaturation("anthropic", "anth-conn-2");
|
||
assert.ok(sig, "expected signal");
|
||
const expected = Date.parse("2026-01-01T00:00:30Z");
|
||
assert.equal(sig!.resetAt, expected, `resetAt should be RFC3339 epoch ${expected}, got ${sig!.resetAt}`);
|
||
// fully exhausted → saturation 1.
|
||
assert.equal(sig!.saturation, 1);
|
||
});
|
||
|
||
test("Anthropic input/output token variants are captured when base tokens header absent", () => {
|
||
_clearRateLimitHeaders();
|
||
// No base anthropic-ratelimit-tokens-*; only the input variant present.
|
||
storeRateLimitHeaders("anth-conn-3", "anthropic", {
|
||
"anthropic-ratelimit-input-tokens-limit": "2000",
|
||
"anthropic-ratelimit-input-tokens-remaining": "500",
|
||
"anthropic-ratelimit-input-tokens-reset": "2026-01-01T00:01:00Z",
|
||
});
|
||
|
||
const sig = getTokenHeaderSaturation("anthropic", "anth-conn-3");
|
||
assert.ok(sig, "expected signal from input-tokens variant");
|
||
// 2000 limit, 500 remaining → 1500 used → 0.75.
|
||
assert.ok(Math.abs(sig!.saturation - 0.75) < 1e-9, `expected ≈0.75, got ${sig!.saturation}`);
|
||
});
|
||
|
||
// ─── OpenAI token headers (reset = DURATION) ─────────────────────────────────
|
||
|
||
test("OpenAI token headers → saturation = 1 − remaining/limit", () => {
|
||
_clearRateLimitHeaders();
|
||
// 90000 limit, 9000 remaining → 81000 used → 0.9.
|
||
storeRateLimitHeaders("oai-conn-1", "openai", {
|
||
"x-ratelimit-limit-tokens": "90000",
|
||
"x-ratelimit-remaining-tokens": "9000",
|
||
"x-ratelimit-reset-tokens": "6m0s",
|
||
});
|
||
|
||
const sig = getTokenHeaderSaturation("openai", "oai-conn-1");
|
||
assert.ok(sig, "expected token-header signal");
|
||
assert.ok(Math.abs(sig!.saturation - 0.9) < 1e-9, `expected ≈0.9, got ${sig!.saturation}`);
|
||
});
|
||
|
||
test('OpenAI reset duration "6m0s" → now + 360s (epoch ms)', () => {
|
||
_clearRateLimitHeaders();
|
||
const before = Date.now();
|
||
storeRateLimitHeaders("oai-conn-2", "openai", {
|
||
"x-ratelimit-limit-tokens": "1000",
|
||
"x-ratelimit-remaining-tokens": "100",
|
||
"x-ratelimit-reset-tokens": "6m0s",
|
||
});
|
||
const after = Date.now();
|
||
|
||
const sig = getTokenHeaderSaturation("openai", "oai-conn-2");
|
||
assert.ok(sig, "expected signal");
|
||
// 6m0s = 360_000 ms in the future, measured from the store time.
|
||
assert.ok(
|
||
sig!.resetAt >= before + 360_000 && sig!.resetAt <= after + 360_000,
|
||
`resetAt should be ≈ now+360000, got ${sig!.resetAt} (window ${before + 360_000}..${after + 360_000})`
|
||
);
|
||
});
|
||
|
||
test('OpenAI reset duration "1s" → now + 1s (epoch ms)', () => {
|
||
_clearRateLimitHeaders();
|
||
const before = Date.now();
|
||
storeRateLimitHeaders("oai-conn-3", "openai", {
|
||
"x-ratelimit-limit-tokens": "1000",
|
||
"x-ratelimit-remaining-tokens": "500",
|
||
"x-ratelimit-reset-tokens": "1s",
|
||
});
|
||
const after = Date.now();
|
||
|
||
const sig = getTokenHeaderSaturation("openai", "oai-conn-3");
|
||
assert.ok(sig, "expected signal");
|
||
assert.ok(
|
||
sig!.resetAt >= before + 1000 && sig!.resetAt <= after + 1000,
|
||
`resetAt should be ≈ now+1000, got ${sig!.resetAt}`
|
||
);
|
||
});
|
||
|
||
test('OpenAI compound duration "1h30m15s" → 5415s', () => {
|
||
_clearRateLimitHeaders();
|
||
const before = Date.now();
|
||
storeRateLimitHeaders("oai-conn-4", "openai", {
|
||
"x-ratelimit-limit-tokens": "1000",
|
||
"x-ratelimit-remaining-tokens": "10",
|
||
"x-ratelimit-reset-tokens": "1h30m15s",
|
||
});
|
||
const after = Date.now();
|
||
|
||
const sig = getTokenHeaderSaturation("openai", "oai-conn-4");
|
||
assert.ok(sig, "expected signal");
|
||
const secs = (1 * 3600 + 30 * 60 + 15) * 1000; // 5_415_000
|
||
assert.ok(
|
||
sig!.resetAt >= before + secs && sig!.resetAt <= after + secs,
|
||
`resetAt should be ≈ now+${secs}, got ${sig!.resetAt}`
|
||
);
|
||
});
|
||
|
||
test('OpenAI fractional-second duration "1.5s" → 1500ms', () => {
|
||
_clearRateLimitHeaders();
|
||
const before = Date.now();
|
||
storeRateLimitHeaders("oai-conn-5", "openai", {
|
||
"x-ratelimit-limit-tokens": "1000",
|
||
"x-ratelimit-remaining-tokens": "10",
|
||
"x-ratelimit-reset-tokens": "1.5s",
|
||
});
|
||
const after = Date.now();
|
||
|
||
const sig = getTokenHeaderSaturation("openai", "oai-conn-5");
|
||
assert.ok(sig, "expected signal");
|
||
assert.ok(
|
||
sig!.resetAt >= before + 1500 && sig!.resetAt <= after + 1500,
|
||
`resetAt should be ≈ now+1500, got ${sig!.resetAt}`
|
||
);
|
||
});
|
||
|
||
// ─── Clamp / guards ──────────────────────────────────────────────────────────
|
||
|
||
test("saturation is clamped to [0,1] (remaining > limit → 0)", () => {
|
||
_clearRateLimitHeaders();
|
||
storeRateLimitHeaders("oai-conn-clamp", "openai", {
|
||
"x-ratelimit-limit-tokens": "1000",
|
||
"x-ratelimit-remaining-tokens": "5000", // nonsensical but must not go negative
|
||
});
|
||
|
||
const sig = getTokenHeaderSaturation("openai", "oai-conn-clamp");
|
||
assert.ok(sig, "expected signal");
|
||
assert.equal(sig!.saturation, 0);
|
||
});
|
||
|
||
test("missing/invalid token headers → no signal (null)", () => {
|
||
_clearRateLimitHeaders();
|
||
storeRateLimitHeaders("oai-conn-none", "openai", {
|
||
"content-type": "application/json",
|
||
});
|
||
assert.equal(getTokenHeaderSaturation("openai", "oai-conn-none"), null);
|
||
|
||
// limit=0 must not divide-by-zero / produce a signal.
|
||
storeRateLimitHeaders("oai-conn-zero", "openai", {
|
||
"x-ratelimit-limit-tokens": "0",
|
||
"x-ratelimit-remaining-tokens": "0",
|
||
});
|
||
assert.equal(getTokenHeaderSaturation("openai", "oai-conn-zero"), null);
|
||
});
|
||
|
||
test("reset is optional — signal still produced without a reset header", () => {
|
||
_clearRateLimitHeaders();
|
||
storeRateLimitHeaders("oai-conn-noreset", "openai", {
|
||
"x-ratelimit-limit-tokens": "1000",
|
||
"x-ratelimit-remaining-tokens": "250",
|
||
});
|
||
const sig = getTokenHeaderSaturation("openai", "oai-conn-noreset");
|
||
assert.ok(sig, "expected signal without reset");
|
||
assert.ok(Math.abs(sig!.saturation - 0.75) < 1e-9);
|
||
assert.equal(sig!.resetAt, null);
|
||
});
|
||
|
||
// ─── No regression on the existing REQUEST path ──────────────────────────────
|
||
|
||
test("REQUEST headers still drive getSaturation (anthropic header fallback unchanged)", async () => {
|
||
_clearSaturationCache();
|
||
_clearRateLimitHeaders();
|
||
// Only request headers (no oauth/usage, no token headers): the legacy fallback
|
||
// must still compute (limit-remaining)/limit = 0.7 for anthropic.
|
||
storeRateLimitHeaders("req-conn", "anthropic", {
|
||
"anthropic-ratelimit-requests-limit": "100",
|
||
"anthropic-ratelimit-requests-remaining": "30",
|
||
});
|
||
|
||
// No oauth token on the connection → fetchAnthropicSaturation falls to the header path.
|
||
satMod.__setAnthropicSaturationDepsForTests({
|
||
loadConnection: async () => ({ id: "req-conn", provider: "anthropic", authType: "apikey" }),
|
||
fetchUsage: async () => ({ message: "no plan window" }),
|
||
});
|
||
|
||
const val = await getSaturation("req-conn", "anthropic", { unit: "requests", window: "hourly" });
|
||
satMod.__setAnthropicSaturationDepsForTests(null);
|
||
assert.ok(Math.abs(val - 0.7) < 1e-9, `expected request-header fallback ≈0.7, got ${val}`);
|
||
});
|
||
|
||
// ─── Generic provider: token-header signal surfaces via getSaturation ────────
|
||
|
||
test("getSaturation(openai) surfaces token-header saturation when usage returns nothing", async () => {
|
||
_clearSaturationCache();
|
||
_clearRateLimitHeaders();
|
||
// usage returns nothing usable → 0; token headers should provide the signal.
|
||
__setGenericUsageFetcherForTests(async () => null);
|
||
storeRateLimitHeaders("oai-gen-1", "openai", {
|
||
"x-ratelimit-limit-tokens": "1000",
|
||
"x-ratelimit-remaining-tokens": "100",
|
||
"x-ratelimit-reset-tokens": "30s",
|
||
});
|
||
|
||
const val = await getSaturation("oai-gen-1", "openai", { unit: "tokens", window: "hourly" });
|
||
assert.ok(Math.abs(val - 0.9) < 1e-9, `expected token-header ≈0.9, got ${val}`);
|
||
});
|
||
|
||
test("getSaturation(openai) prefers real usage percent over token headers when usage is present", async () => {
|
||
_clearSaturationCache();
|
||
_clearRateLimitHeaders();
|
||
// usage reports 0.5 (50% used) → authoritative; token headers say 0.9 but must not override.
|
||
__setGenericUsageFetcherForTests(async () => ({ percentUsed: 0.5 }));
|
||
storeRateLimitHeaders("oai-gen-2", "openai", {
|
||
"x-ratelimit-limit-tokens": "1000",
|
||
"x-ratelimit-remaining-tokens": "100", // → 0.9 if it leaked through
|
||
});
|
||
|
||
const val = await getSaturation("oai-gen-2", "openai", { unit: "tokens", window: "hourly" });
|
||
assert.ok(Math.abs(val - 0.5) < 1e-9, `expected usage ≈0.5 to win, got ${val}`);
|
||
});
|
||
|
||
test("getSaturation(openai) fails open to 0 when neither usage nor token headers exist", async () => {
|
||
_clearSaturationCache();
|
||
_clearRateLimitHeaders();
|
||
__setGenericUsageFetcherForTests(async () => null);
|
||
|
||
const val = await getSaturation("oai-gen-3", "openai", { unit: "tokens", window: "hourly" });
|
||
assert.equal(val, 0);
|
||
});
|
||
|
||
// ─── Cross-key / cross-provider isolation ────────────────────────────────────
|
||
|
||
test("token-header signal is keyed by (provider, connectionId)", () => {
|
||
_clearRateLimitHeaders();
|
||
storeRateLimitHeaders("shared-id", "openai", {
|
||
"x-ratelimit-limit-tokens": "1000",
|
||
"x-ratelimit-remaining-tokens": "100", // 0.9
|
||
});
|
||
storeRateLimitHeaders("shared-id", "anthropic", {
|
||
"anthropic-ratelimit-tokens-limit": "1000",
|
||
"anthropic-ratelimit-tokens-remaining": "900", // 0.1
|
||
});
|
||
|
||
const oai = getTokenHeaderSaturation("openai", "shared-id");
|
||
const anth = getTokenHeaderSaturation("anthropic", "shared-id");
|
||
assert.ok(oai && Math.abs(oai.saturation - 0.9) < 1e-9);
|
||
assert.ok(anth && Math.abs(anth.saturation - 0.1) < 1e-9);
|
||
// A provider with no token headers stored → null.
|
||
assert.equal(getTokenHeaderSaturation("openai", "never-seen"), null);
|
||
});
|