From 43335f078992658b58ea094c1e2c1d95998460a1 Mon Sep 17 00:00:00 2001 From: benzntech Date: Thu, 13 Aug 2026 09:26:12 -0300 Subject: [PATCH] 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> --- open-sse/services/combo.ts | 15 +---- src/app/api/settings/quota/state/route.ts | 77 +++++++++++++++-------- src/lib/quota/quotaAdapters.ts | 5 +- tests/unit/quota-phase2.test.ts | 58 ++++++++++++++--- 4 files changed, 107 insertions(+), 48 deletions(-) diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 4b126df7f9..76e04d9b17 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -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 diff --git a/src/app/api/settings/quota/state/route.ts b/src/app/api/settings/quota/state/route.ts index 7b58e168fd..a47128afed 100644 --- a/src/app/api/settings/quota/state/route.ts +++ b/src/app/api/settings/quota/state/route.ts @@ -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, + }); } } diff --git a/src/lib/quota/quotaAdapters.ts b/src/lib/quota/quotaAdapters.ts index e202112f5c..1bdb636451 100644 --- a/src/lib/quota/quotaAdapters.ts +++ b/src/lib/quota/quotaAdapters.ts @@ -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), + }); } diff --git a/tests/unit/quota-phase2.test.ts b/tests/unit/quota-phase2.test.ts index ee60ea3d62..64e0600e2a 100644 --- a/tests/unit/quota-phase2.test.ts +++ b/tests/unit/quota-phase2.test.ts @@ -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);