mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-18 21:22:28 +03:00
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:
@@ -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
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user