Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 43 additions & 4 deletions server/src/__tests__/lib/client-context.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { afterEach, describe, expect, it } from 'vitest';
import type { NextFunction, Request, Response } from 'express';
import { clientContextMiddleware, getClientContext } from '../../lib/client-context.js';
import { clientContextMiddleware, getClientContext, setObservedRequestTokens, getObservedRequestTokens } from '../../lib/client-context.js';

// Minimal fake req: the middleware only touches headers, socket, and app.
function fakeReq(headers: Record<string, string | string[]>, remoteAddress?: string, trustProxy?: boolean): Request {
Expand All @@ -22,7 +22,7 @@ describe('clientContextMiddleware', () => {

it('captures the socket peer address and user agent', () => {
const ctx = contextFor(fakeReq({ 'user-agent': 'curl/8.6.0' }, '192.168.0.42'));
expect(ctx).toEqual({ ip: '192.168.0.42', userAgent: 'curl/8.6.0', agent: 'unknown' });
expect(ctx).toEqual({ ip: '192.168.0.42', userAgent: 'curl/8.6.0', agent: 'unknown', observedRequestTokens: null });
});

it('prefers the first X-Forwarded-For hop when trust proxy is enabled', () => {
Expand All @@ -47,10 +47,49 @@ describe('clientContextMiddleware', () => {
it('stores nulls when REQUEST_ANALYTICS_LOG_CLIENT=false', () => {
process.env.REQUEST_ANALYTICS_LOG_CLIENT = 'false';
const ctx = contextFor(fakeReq({ 'user-agent': 'curl/8.6.0' }, '192.168.0.42'));
expect(ctx).toEqual({ ip: null, userAgent: null, agent: null });
expect(ctx).toEqual({ ip: null, userAgent: null, agent: null, observedRequestTokens: null });
});

it('returns nulls outside any request scope', () => {
expect(getClientContext()).toEqual({ ip: null, userAgent: null, agent: null });
expect(getClientContext()).toEqual({ ip: null, userAgent: null, agent: null, observedRequestTokens: null });
});
});

describe('setObservedRequestTokens / getObservedRequestTokens', () => {
it('returns null outside any request scope and is a no-op setter', () => {
expect(getObservedRequestTokens()).toBeNull();
// Setter is a no-op when no store is active — must not throw.
setObservedRequestTokens(12345);
expect(getObservedRequestTokens()).toBeNull();
});

it('sticky-writes the value inside a request scope', () => {
let seen: number | null = null;
clientContextMiddleware({ headers: {}, socket: {} } as unknown as Request, {} as Response, (() => {
setObservedRequestTokens(36_532);
seen = getObservedRequestTokens();
}) as NextFunction);
expect(seen).toBe(36_532);
});

it('never decreases — max of current and incoming wins', () => {
clientContextMiddleware({ headers: {}, socket: {} } as unknown as Request, {} as Response, (() => {
setObservedRequestTokens(40_000);
setObservedRequestTokens(10_000); // smaller, must be ignored
setObservedRequestTokens(50_000); // larger, takes over
expect(getObservedRequestTokens()).toBe(50_000);
}) as NextFunction);
});

it('scope is per-request — one request cannot leak into another', () => {
const captured: Array<number | null> = [];
clientContextMiddleware({ headers: {}, socket: {} } as unknown as Request, {} as Response, (() => {
setObservedRequestTokens(99_999);
captured.push(getObservedRequestTokens());
}) as NextFunction);
clientContextMiddleware({ headers: {}, socket: {} } as unknown as Request, {} as Response, (() => {
captured.push(getObservedRequestTokens()); // new scope → null
}) as NextFunction);
expect(captured).toEqual([99_999, null]);
});
});
102 changes: 102 additions & 0 deletions server/src/__tests__/lib/provider-size-parser.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
import { describe, expect, it } from 'vitest';
import { parseProviderReportedSize } from '../../lib/provider-size-parser.js';

// Bodies were captured verbatim from the live requests table (provider + status
// = error, last 30 days) and trimmed only to fit. Each test pins one real
// observed error message shape so a provider changing its wording trips the
// test and forces a parser update.
describe('parseProviderReportedSize', () => {
describe('groq', () => {
it('extracts Requested N from a TPM 413', () => {
const msg = 'Groq API error 413: Request too large for model `openai/gpt-oss-120b` in organization `org_01kptjck7bejta1btzc11cccy2` service tier `on_demand` on tokens per minute (TPM): Limit 8000, Requested 36532, please reduce your message size and try again.';
expect(parseProviderReportedSize('groq', msg)).toBe(36532);
});

it('extracts Requested N from llama-3.1-8b-instant 413', () => {
const msg = 'Groq API error 413: Request too large for model `llama-3.1-8b-instant` in organization `org_01kptjck7bejta1btzc11cccy2` service tier `on_demand` on tokens per minute (TPM): Limit 6000, Requested 36783, please reduce your message size and try again.';
expect(parseProviderReportedSize('groq', msg)).toBe(36783);
});

it('returns null for the bare "Request Entity Too Large" body', () => {
// This shape appears 627 times in 30d but carries no number. Returning
// a number here would falsely report a request size.
expect(parseProviderReportedSize('groq', 'Groq API error 413: Request Entity Too Large')).toBeNull();
});

it('handles thousand-separator commas', () => {
const msg = 'Groq API error 413: ... Limit 8000, Requested 36,532, please reduce ...';
expect(parseProviderReportedSize('groq', msg)).toBe(36532);
});
});

describe('openrouter', () => {
it('extracts the total from "requested about N tokens (...)"', () => {
const msg = 'OpenRouter API error 400: This endpoint\'s maximum context length is 65536 tokens. However, you requested about 68982 tokens (4982 of text input, 64000 in the output). Please reduce the length of either one, or use the context-compression.';
expect(parseProviderReportedSize('openrouter', msg)).toBe(68982);
});

it('handles variation in thousand separators and whitespace', () => {
const msg = 'OpenRouter API error 400: requested about 68,847 tokens (4847 of text input, 64000 in the output).';
expect(parseProviderReportedSize('openrouter', msg)).toBe(68847);
});

it('returns null on a non-size error', () => {
expect(parseProviderReportedSize('openrouter', 'OpenRouter API error 429: Provider returned error')).toBeNull();
});
});

describe('cloudflare', () => {
it('prefers the input-only number from a 400 context-length body', () => {
const msg = 'Cloudflare API error 400: AiError: AiError: {"error":{"message":"This model\'s maximum context length is 24000 tokens. However, you requested 256 output tokens and your prompt contains at least 23745 input tokens, for a total of at least 24001 tokens."}}';
expect(parseProviderReportedSize('cloudflare', msg)).toBe(23745);
});

it('falls back to the combined total from a 413 "tokens (N) exceeded" body', () => {
const msg = 'Cloudflare API error 413: AiError: Ai: The estimated number of input and maximum output tokens (24092) exceeded this model context window limit (24000). (1ffb6b51-7168-4e29-a4ab-378d87917a79)';
expect(parseProviderReportedSize('cloudflare', msg)).toBe(24092);
});

it('returns null on a non-size error', () => {
expect(parseProviderReportedSize('cloudflare', 'Cloudflare API error 429: AiError: you have used up your daily free allocation')).toBeNull();
});
});

describe('github', () => {
it('returns null even though the body is parseable (limit only, not request size)', () => {
// "Max size: 8000 tokens" is the LIMIT ceiling. Returning 8000 would
// cause every subsequent model with TPM < 8000 to be skipped for the
// rest of the request — wildly wrong. The parser must refuse.
const msg = 'GitHub Models API error 413: Request body too large for gpt-4.1 model. Max size: 8000 tokens.';
expect(parseProviderReportedSize('github', msg)).toBeNull();
});
});

describe('providers without a parser', () => {
it('returns null for ollama, nvidia, anthropic, google, opencode, llm7, cerebras, custom', () => {
for (const p of ['ollama', 'nvidia', 'anthropic', 'google', 'opencode', 'llm7', 'cerebras', 'custom']) {
expect(parseProviderReportedSize(p, 'some error containing 12345 tokens')).toBeNull();
}
});
});

describe('edge cases', () => {
it('returns null for empty / nullish message', () => {
expect(parseProviderReportedSize('groq', undefined)).toBeNull();
expect(parseProviderReportedSize('groq', null)).toBeNull();
expect(parseProviderReportedSize('groq', '')).toBeNull();
});

it('returns null when the number is missing or malformed', () => {
expect(parseProviderReportedSize('groq', 'Limit zero, Requested none, please reduce')).toBeNull();
expect(parseProviderReportedSize('openrouter', 'requested about zero tokens')).toBeNull();
});

it('ignores zero or negative numbers', () => {
// Defensive — shouldn't happen with real providers, but a stray "0" in
// an error template would otherwise set observedRequestTokens=0 and
// silently disable the gate for the rest of the request.
const msg = 'Groq API error 413: ... Limit 8000, Requested 0, ...';
expect(parseProviderReportedSize('groq', msg)).toBeNull();
});
});
});
157 changes: 157 additions & 0 deletions server/src/__tests__/services/router.test.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { describe, it, expect, beforeAll, beforeEach, afterEach, vi } from 'vitest';
import type { NextFunction, Request, Response } from 'express';
import { initDb, getDb } from '../../db/index.js';
import { encrypt } from '../../lib/crypto.js';
import {
Expand All @@ -8,6 +9,7 @@ import {
setRoutingStrategy,
} from '../../services/router.js';
import { setCooldown } from '../../services/ratelimit.js';
import { clientContextMiddleware, setObservedRequestTokens } from '../../lib/client-context.js';

describe('Router', () => {
beforeAll(() => {
Expand Down Expand Up @@ -276,3 +278,158 @@ describe('Router exhaustion diagnostics (issue _1)', () => {
expect(caught.diagnostics.some((d: string) => /cooldown/.test(d))).toBe(true);
});
});

// Run a function inside a client-context scope with the given observedRequestTokens.
// routeRequest() reads the AsyncLocalStorage store directly, so without an
// enclosing scope the observed-size gate is a no-op (the production hot path
// is always inside clientContextMiddleware).
function routeInScope<T>(observed: number | null, fn: () => T): T {
const fakeReq = { headers: {}, socket: {} } as unknown as Request;
let result!: T;
clientContextMiddleware(fakeReq, {} as Response, (() => {
if (observed != null) setObservedRequestTokens(observed);
result = fn();
}) as NextFunction);
return result;
}

describe('Router sticky provider-reported size gate', () => {
beforeAll(() => {
process.env.ENCRYPTION_KEY = '0'.repeat(64);
initDb(':memory:');
});

beforeEach(() => {
const db = getDb();
setRoutingStrategy('priority');
db.prepare('DELETE FROM api_keys').run();
db.prepare("DELETE FROM settings WHERE key = 'active_profile_id'").run();
db.prepare('DELETE FROM rate_limit_cooldowns').run();
db.prepare('DELETE FROM rate_limit_usage').run();
// Reset fallback_config priorities to the legacy-baseline order
// (intelligence rank ascending) so tests that ran before us can't
// leave a model at priority 0 and bump the chain shape. Mirrors the
// first describe block's reset so tests are independent.
const models = db.prepare(
'SELECT id FROM models ORDER BY intelligence_rank ASC',
).all() as { id: number }[];
const update = db.prepare(
'UPDATE fallback_config SET priority = ? WHERE model_db_id = ?',
);
for (let i = 0; i < models.length; i++) {
update.run(i + 1, models[i].id);
}
});

afterEach(() => {
vi.useRealTimers();
});

// Helper: pin one model to a specific (tpm_limit, platform) and place it at
// the front of the chain. Returns its model_db_id.
function pinModel(platform: string, modelId: string, tpmLimit: number | null): number {
const db = getDb();
const row = db.prepare(
'SELECT id FROM models WHERE platform = ? AND model_id = ?',
).get(platform, modelId) as { id: number } | undefined;
if (!row) throw new Error(`model ${platform}/${modelId} not seeded`);
db.prepare('UPDATE models SET tpm_limit = ? WHERE id = ?').run(tpmLimit, row.id);
// Use priority=0 (top of chain) and bump every other model way down so
// collisions with seeded priorities can't push another model in front.
db.prepare('UPDATE fallback_config SET priority = 0 WHERE model_db_id = ?').run(row.id);
db.prepare('UPDATE fallback_config SET priority = priority + 10000 WHERE model_db_id != ?').run(row.id);
return row.id;
}

it('skips a model whose tpm_limit cannot fit an observed request size', () => {
const db = getDb();
// Insert keys for TWO groq models with very different TPMs. By adding
// both keys we get to test the model-selection gate without fighting
// fallback_config priority plumbing — only the model whose TPM can't
// fit the observed size should be skipped.
const lowTpmKey = encrypt('test-groq-key-low');
db.prepare(`
INSERT INTO api_keys (platform, label, encrypted_key, iv, auth_tag, status, enabled)
VALUES (?, ?, ?, ?, ?, ?, ?)
`).run('groq', 'low-tpm-key', lowTpmKey.encrypted, lowTpmKey.iv, lowTpmKey.authTag, 'healthy', 1);

const highTpmKey = encrypt('test-groq-key-high');
db.prepare(`
INSERT INTO api_keys (platform, label, encrypted_key, iv, auth_tag, status, enabled)
VALUES (?, ?, ?, ?, ?, ?, ?)
`).run('groq', 'high-tpm-key', highTpmKey.encrypted, highTpmKey.iv, highTpmKey.authTag, 'healthy', 1);

const lowTpmId = pinModel('groq', 'llama-3.1-8b-instant', 6_000);
const highTpmId = pinModel('groq', 'groq/compound', 70_000);
expect(lowTpmId).not.toBe(highTpmId);

// Pin the low-TPM model first in the chain so it's preferred when both
// are eligible, then the high-TPM model as fallback.
db.prepare('UPDATE fallback_config SET priority = 0 WHERE model_db_id = ?').run(lowTpmId);
db.prepare('UPDATE fallback_config SET priority = 1 WHERE model_db_id = ?').run(highTpmId);
db.prepare('UPDATE fallback_config SET priority = priority + 10000 WHERE model_db_id NOT IN (?, ?)').run(lowTpmId, highTpmId);

// Without an observed size, the pinned first model (low TPM) wins.
const lowResult = routeInScope(null, () => routeRequest(500));
expect(lowResult.modelDbId).toBe(lowTpmId);

// With observed=36532 set on this request, the 6K-TPM model is skipped
// pre-flight and the 70K-TPM model wins.
const highResult = routeInScope(36_532, () => routeRequest(500));
expect(highResult.modelDbId).toBe(highTpmId);
});

it('records the new skip reason in the exhaustion diagnostic when ALL models are too small', () => {
const db = getDb();
const groqKey = encrypt('test-groq-key');
db.prepare(`
INSERT INTO api_keys (platform, label, encrypted_key, iv, auth_tag, status, enabled)
VALUES (?, ?, ?, ?, ?, ?, ?)
`).run('groq', 'test', groqKey.encrypted, groqKey.iv, groqKey.authTag, 'healthy', 1);

// Configure every Groq model to a TPM below the observed size, so the
// gate rejects all of them.
db.prepare("UPDATE models SET tpm_limit = 1000 WHERE platform = 'groq'").run();
db.prepare("UPDATE models SET tpd_limit = NULL, context_window = 10000000 WHERE platform = 'groq'").run();

let caught: any;
try {
routeInScope(36_532, () => routeRequest(500));
} catch (e) {
caught = e;
}
expect(caught).toBeDefined();
expect(Array.isArray(caught.diagnostics)).toBe(true);
expect(caught.diagnostics.some((d: string) => /request-too-large-for-tpm/.test(d))).toBe(true);
});

it('does NOT skip when tpm_limit is NULL (unknown ceiling — fall through to the local estimator)', () => {
const db = getDb();
const groqKey = encrypt('test-groq-key');
db.prepare(`
INSERT INTO api_keys (platform, label, encrypted_key, iv, auth_tag, status, enabled)
VALUES (?, ?, ?, ?, ?, ?, ?)
`).run('groq', 'test', groqKey.encrypted, groqKey.iv, groqKey.authTag, 'healthy', 1);

const id = pinModel('groq', 'llama-3.1-8b-instant', null);
db.prepare("UPDATE models SET tpm_limit = NULL WHERE id = ?").run(id);

// Unknown ceiling: even with a huge observed size, the gate does not
// reject — the existing canUseTokens headroom check still applies, but
// a null tpm_limit means no pre-flight rejection.
expect(routeInScope(1_000_000, () => routeRequest(500).modelDbId)).toBe(id);
});

it('does NOT skip when observed size fits within tpm_limit', () => {
const db = getDb();
const groqKey = encrypt('test-groq-key');
db.prepare(`
INSERT INTO api_keys (platform, label, encrypted_key, iv, auth_tag, status, enabled)
VALUES (?, ?, ?, ?, ?, ?, ?)
`).run('groq', 'test', groqKey.encrypted, groqKey.iv, groqKey.authTag, 'healthy', 1);

// 8K TPM, 5K observed — fits.
const id = pinModel('groq', 'llama-3.1-8b-instant', 8_000);
expect(routeInScope(5_000, () => routeRequest(500).modelDbId)).toBe(id);
});
});
Loading
Loading