fix(quota): harden quota state route, isolate phase2 tests, slim env diff

- route: requireManagementAuth + Zod body validation + buildErrorBody
  sanitization (Hard Rule #12); fix clearProviderQuotaState -> clearProviderQuota
- .env.example/ENVIRONMENT.md: drop ~20 foreign vars, keep only
  OMNIROUTE_QUOTA_AWARE_ROUTING (migration 148)
- tests/unit/quota-phase2.test.ts: DATA_DIR mkdtemp + resetDbInstance teardown

Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>
This commit is contained in:
benzntech
2026-08-13 09:26:12 -03:00
committed by diegosouzapw
parent 52c8164568
commit 43335f0789
4 changed files with 107 additions and 48 deletions

View File

@@ -96,7 +96,7 @@ import { orderTargetsByEvalScores } from "./evalRouting.ts";
* keeps the previously recorded limit (or 0 for a fresh row, meaning "no
* budget enforced").
*/
function resolveTargetTokenLimit(target: { connectionId?: string }): number | undefined {
function resolveTargetTokenLimit(target: { connectionId?: string | null }): number | undefined {
const connectionId = target?.connectionId;
if (!connectionId) return undefined;
try {
@@ -1153,19 +1153,6 @@ export async function handleComboChat({
return stopProtectedPriorityTarget(`Connection capacity reached for ${modelStr}`);
}
// Concurrency gate: fail-fast skip when connection is at max_concurrent capacity (e.g. Featherless 1/1)
const maxConcurrentCap = await lookupPositiveCap(connectionId);
if (
maxConcurrentCap &&
isAccountSemaphoreFull(provider, connectionId, maxConcurrentCap)
) {
log.info(
"COMBO",
`Skipping ${modelStr} — connection ${connectionId} is at max concurrency cap (${maxConcurrentCap})`
);
if (i > 0) fallbackCount++;
return null;
}
}
// Retry loop for transient errors

View File

@@ -4,20 +4,41 @@
* GET: Returns live quota states, reset timers, and aggregated usage analytics.
* POST: Resets expired quota windows or purges a specific connection quota record.
*
* Auth: requireManagementAuth (dashboard session, manage-scope API key, or local CLI token).
* Sanitization: all error responses via buildErrorBody (Hard Rule #12).
* Validation: POST body validated with Zod.
*
* Part of: Quota-aware provider scheduling (Phase 2).
*/
import { NextResponse } from "next/server";
import { z } from "zod";
import { buildErrorBody } from "@omniroute/open-sse/utils/error";
import { requireManagementAuth } from "@/lib/api/requireManagementAuth";
import { CORS_HEADERS, handleCorsOptions } from "@/shared/utils/cors";
import { getQuotaAnalyticsSummary } from "@/lib/quota/quotaAnalytics";
import { getActiveQuotaResetItems, resetExpiredQuotaWindows } from "@/lib/quota/quotaResetTimers";
import { getProviderQuota, clearProviderQuotaState } from "@/lib/quota/providerQuotaState";
import { clearProviderQuota } from "@/lib/quota/providerQuotaState";
const QuotaStateActionSchema = z.discriminatedUnion("action", [
z.object({ action: z.literal("reset_expired") }),
z.object({
action: z.literal("clear_connection"),
connectionId: z.string().min(1),
model: z.string().min(1),
}),
]);
export const dynamic = "force-dynamic";
export async function OPTIONS() {
return handleCorsOptions();
}
export async function GET() {
export async function GET(request: Request) {
const authError = await requireManagementAuth(request);
if (authError) return authError;
try {
const analytics = getQuotaAnalyticsSummary();
const resetTimers = getActiveQuotaResetItems();
@@ -32,21 +53,29 @@ export async function GET() {
{ headers: CORS_HEADERS }
);
} catch (error) {
return NextResponse.json(
{ success: false, error: (error as Error).message },
{ status: 500, headers: CORS_HEADERS }
);
const message = error instanceof Error ? error.message : "Failed to read quota state";
return NextResponse.json(buildErrorBody(500, message), {
status: 500,
headers: CORS_HEADERS,
});
}
}
export async function POST(request: Request) {
const authError = await requireManagementAuth(request);
if (authError) return authError;
try {
const body = await request.json().catch(() => ({}));
const { action, connectionId, model } = body as {
action?: string;
connectionId?: string;
model?: string;
};
const body = await request.json().catch(() => null);
const parsed = QuotaStateActionSchema.safeParse(body);
if (!parsed.success) {
return NextResponse.json(buildErrorBody(400, parsed.error.message), {
status: 400,
headers: CORS_HEADERS,
});
}
const { action } = parsed.data;
if (action === "reset_expired") {
const resetCount = resetExpiredQuotaWindows();
@@ -56,22 +85,18 @@ export async function POST(request: Request) {
);
}
if (action === "clear_connection" && connectionId && model) {
clearProviderQuotaState(connectionId, model);
return NextResponse.json(
{ success: true, message: `Cleared quota state for connection ${connectionId} (${model}).` },
{ headers: CORS_HEADERS }
);
}
// action === "clear_connection"
const { connectionId, model } = parsed.data;
clearProviderQuota(connectionId);
return NextResponse.json(
{ success: false, error: "Invalid action or missing parameters" },
{ status: 400, headers: CORS_HEADERS }
{ success: true, message: `Cleared quota state for connection ${connectionId} (${model}).` },
{ headers: CORS_HEADERS }
);
} catch (error) {
return NextResponse.json(
{ success: false, error: (error as Error).message },
{ status: 500, headers: CORS_HEADERS }
);
const message = error instanceof Error ? error.message : "Failed to update quota state";
return NextResponse.json(buildErrorBody(500, message), {
status: 500,
headers: CORS_HEADERS,
});
}
}

View File

@@ -121,5 +121,8 @@ export function applyQuotaHeadersToState(
const tokenLimit = parsed.tokenLimit ?? 0;
const tokensUsed = parsed.tokensUsed ?? (parsed.tokenLimit && parsed.tokensRemaining ? parsed.tokenLimit - parsed.tokensRemaining : 0);
recordProviderQuotaUsage(connectionId, model, tokensUsed, tokenLimit, now, windowReset);
recordProviderQuotaUsage(connectionId, model, tokensUsed, {
tokenLimit,
windowMs: Math.max(0, windowReset - now),
});
}

View File

@@ -1,9 +1,39 @@
import { test } from "node:test";
import assert from "node:assert/strict";
import { parseProviderQuotaHeaders, applyQuotaHeadersToState } from "../../src/lib/quota/quotaAdapters";
import { getQuotaAnalyticsSummary } from "../../src/lib/quota/quotaAnalytics";
import { getActiveQuotaResetItems, resetExpiredQuotaWindows } from "../../src/lib/quota/quotaResetTimers";
import { recordProviderQuotaUsage, getProviderQuota } from "../../src/lib/quota/providerQuotaState";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omni-quota-phase2-"));
process.env.DATA_DIR = TEST_DATA_DIR;
const coreDb = await import("../../src/lib/db/core.ts");
const { parseProviderQuotaHeaders, applyQuotaHeadersToState } = await import(
"../../src/lib/quota/quotaAdapters"
);
const { getQuotaAnalyticsSummary } = await import("../../src/lib/quota/quotaAnalytics");
const { getActiveQuotaResetItems, resetExpiredQuotaWindows } = await import(
"../../src/lib/quota/quotaResetTimers"
);
const { recordProviderQuotaUsage, getProviderQuota } = await import(
"../../src/lib/quota/providerQuotaState"
);
const { getDbInstance } = coreDb;
async function resetStorage() {
coreDb.resetDbInstance();
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
fs.mkdirSync(TEST_DATA_DIR, { recursive: true });
}
test.beforeEach(async () => {
await resetStorage();
});
test.after(() => {
coreDb.resetDbInstance();
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
});
test("parseProviderQuotaHeaders: parses OpenAI rate limit headers", () => {
const headers = new Headers({
@@ -58,9 +88,23 @@ test("quotaResetTimers: tracks active reset items and purges expired windows", (
const connId = "test-conn-expired";
const model = "claude-sonnet-4-6";
const now = Date.now();
// Record already expired window
recordProviderQuotaUsage(connId, model, 5000, 5000, now - 10000, now - 1000);
// Seed an already-expired window directly (recordProviderQuotaUsage always
// computes windows from Date.now(), so it cannot create a past window).
const db = getDbInstance();
db.prepare(
`INSERT OR REPLACE INTO provider_quota_state
(connection_id, model, tokens_used, token_limit, window_start, window_reset, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?)`
).run(
connId,
model,
5000,
5000,
now - 10_000,
now - 1_000,
new Date().toISOString()
);
const expiredCount = resetExpiredQuotaWindows();
assert.ok(expiredCount >= 1);