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
|
* keeps the previously recorded limit (or 0 for a fresh row, meaning "no
|
||||||
* budget enforced").
|
* budget enforced").
|
||||||
*/
|
*/
|
||||||
function resolveTargetTokenLimit(target: { connectionId?: string }): number | undefined {
|
function resolveTargetTokenLimit(target: { connectionId?: string | null }): number | undefined {
|
||||||
const connectionId = target?.connectionId;
|
const connectionId = target?.connectionId;
|
||||||
if (!connectionId) return undefined;
|
if (!connectionId) return undefined;
|
||||||
try {
|
try {
|
||||||
@@ -1153,19 +1153,6 @@ export async function handleComboChat({
|
|||||||
return stopProtectedPriorityTarget(`Connection capacity reached for ${modelStr}`);
|
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
|
// Retry loop for transient errors
|
||||||
|
|||||||
@@ -4,20 +4,41 @@
|
|||||||
* GET: Returns live quota states, reset timers, and aggregated usage analytics.
|
* GET: Returns live quota states, reset timers, and aggregated usage analytics.
|
||||||
* POST: Resets expired quota windows or purges a specific connection quota record.
|
* 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).
|
* Part of: Quota-aware provider scheduling (Phase 2).
|
||||||
*/
|
*/
|
||||||
|
|
||||||
import { NextResponse } from "next/server";
|
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 { CORS_HEADERS, handleCorsOptions } from "@/shared/utils/cors";
|
||||||
import { getQuotaAnalyticsSummary } from "@/lib/quota/quotaAnalytics";
|
import { getQuotaAnalyticsSummary } from "@/lib/quota/quotaAnalytics";
|
||||||
import { getActiveQuotaResetItems, resetExpiredQuotaWindows } from "@/lib/quota/quotaResetTimers";
|
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() {
|
export async function OPTIONS() {
|
||||||
return handleCorsOptions();
|
return handleCorsOptions();
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function GET() {
|
export async function GET(request: Request) {
|
||||||
|
const authError = await requireManagementAuth(request);
|
||||||
|
if (authError) return authError;
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const analytics = getQuotaAnalyticsSummary();
|
const analytics = getQuotaAnalyticsSummary();
|
||||||
const resetTimers = getActiveQuotaResetItems();
|
const resetTimers = getActiveQuotaResetItems();
|
||||||
@@ -32,21 +53,29 @@ export async function GET() {
|
|||||||
{ headers: CORS_HEADERS }
|
{ headers: CORS_HEADERS }
|
||||||
);
|
);
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
return NextResponse.json(
|
const message = error instanceof Error ? error.message : "Failed to read quota state";
|
||||||
{ success: false, error: (error as Error).message },
|
return NextResponse.json(buildErrorBody(500, message), {
|
||||||
{ status: 500, headers: CORS_HEADERS }
|
status: 500,
|
||||||
);
|
headers: CORS_HEADERS,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function POST(request: Request) {
|
export async function POST(request: Request) {
|
||||||
|
const authError = await requireManagementAuth(request);
|
||||||
|
if (authError) return authError;
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const body = await request.json().catch(() => ({}));
|
const body = await request.json().catch(() => null);
|
||||||
const { action, connectionId, model } = body as {
|
const parsed = QuotaStateActionSchema.safeParse(body);
|
||||||
action?: string;
|
if (!parsed.success) {
|
||||||
connectionId?: string;
|
return NextResponse.json(buildErrorBody(400, parsed.error.message), {
|
||||||
model?: string;
|
status: 400,
|
||||||
};
|
headers: CORS_HEADERS,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
const { action } = parsed.data;
|
||||||
|
|
||||||
if (action === "reset_expired") {
|
if (action === "reset_expired") {
|
||||||
const resetCount = resetExpiredQuotaWindows();
|
const resetCount = resetExpiredQuotaWindows();
|
||||||
@@ -56,22 +85,18 @@ export async function POST(request: Request) {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
if (action === "clear_connection" && connectionId && model) {
|
// action === "clear_connection"
|
||||||
clearProviderQuotaState(connectionId, model);
|
const { connectionId, model } = parsed.data;
|
||||||
return NextResponse.json(
|
clearProviderQuota(connectionId);
|
||||||
{ success: true, message: `Cleared quota state for connection ${connectionId} (${model}).` },
|
|
||||||
{ headers: CORS_HEADERS }
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
return NextResponse.json(
|
return NextResponse.json(
|
||||||
{ success: false, error: "Invalid action or missing parameters" },
|
{ success: true, message: `Cleared quota state for connection ${connectionId} (${model}).` },
|
||||||
{ status: 400, headers: CORS_HEADERS }
|
{ headers: CORS_HEADERS }
|
||||||
);
|
);
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
return NextResponse.json(
|
const message = error instanceof Error ? error.message : "Failed to update quota state";
|
||||||
{ success: false, error: (error as Error).message },
|
return NextResponse.json(buildErrorBody(500, message), {
|
||||||
{ status: 500, headers: CORS_HEADERS }
|
status: 500,
|
||||||
);
|
headers: CORS_HEADERS,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -121,5 +121,8 @@ export function applyQuotaHeadersToState(
|
|||||||
const tokenLimit = parsed.tokenLimit ?? 0;
|
const tokenLimit = parsed.tokenLimit ?? 0;
|
||||||
const tokensUsed = parsed.tokensUsed ?? (parsed.tokenLimit && parsed.tokensRemaining ? parsed.tokenLimit - parsed.tokensRemaining : 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 { test } from "node:test";
|
||||||
import assert from "node:assert/strict";
|
import assert from "node:assert/strict";
|
||||||
import { parseProviderQuotaHeaders, applyQuotaHeadersToState } from "../../src/lib/quota/quotaAdapters";
|
import fs from "node:fs";
|
||||||
import { getQuotaAnalyticsSummary } from "../../src/lib/quota/quotaAnalytics";
|
import os from "node:os";
|
||||||
import { getActiveQuotaResetItems, resetExpiredQuotaWindows } from "../../src/lib/quota/quotaResetTimers";
|
import path from "node:path";
|
||||||
import { recordProviderQuotaUsage, getProviderQuota } from "../../src/lib/quota/providerQuotaState";
|
|
||||||
|
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", () => {
|
test("parseProviderQuotaHeaders: parses OpenAI rate limit headers", () => {
|
||||||
const headers = new 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 connId = "test-conn-expired";
|
||||||
const model = "claude-sonnet-4-6";
|
const model = "claude-sonnet-4-6";
|
||||||
const now = Date.now();
|
const now = Date.now();
|
||||||
|
|
||||||
// Record already expired window
|
// Seed an already-expired window directly (recordProviderQuotaUsage always
|
||||||
recordProviderQuotaUsage(connId, model, 5000, 5000, now - 10000, now - 1000);
|
// 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();
|
const expiredCount = resetExpiredQuotaWindows();
|
||||||
assert.ok(expiredCount >= 1);
|
assert.ok(expiredCount >= 1);
|
||||||
|
|||||||
Reference in New Issue
Block a user