From bd4407cb64a26a35a7890855dd87e2c996d17c41 Mon Sep 17 00:00:00 2001 From: Diego Rodrigues de Sa e Souza Date: Thu, 6 Aug 2026 10:40:20 -0300 Subject: [PATCH] feat(sse): New-API/One-API/Sub2API aggregator balance detection (#9415) (#9539) Validated in local merge-train (diegosouzapw batch) --- .../9415-newapi-sub2api-aggregator-balance.md | 6 + .../services/newApiAggregatorQuotaFetcher.ts | 225 ++++++++++++++ open-sse/services/quotaMonitor.ts | 13 +- open-sse/services/quotaPreflight.ts | 36 ++- .../[id]/components/modals/AddApiKeyModal.tsx | 2 + .../modals/EditCompatibleNodeModal.tsx | 47 ++- .../components/modals/EditConnectionModal.tsx | 8 + .../modals/NewApiAggregatorFields.tsx | 59 ++++ .../modals/connectionProviderSpecificData.ts | 27 ++ .../components/AddCompatibleProviderModal.tsx | 50 +++- src/i18n/messages/en.json | 7 + src/i18n/messages/pt-BR.json | 7 + .../constants/featureFlagDefinitions.ts | 15 +- src/shared/validation/providerSpecificData.ts | 24 ++ ...wapi-aggregator-preflight-dispatch.test.ts | 119 ++++++++ .../newapi-aggregator-quota-fetcher.test.ts | 280 ++++++++++++++++++ .../provider-specific-data-schema.test.ts | 64 ++++ 17 files changed, 982 insertions(+), 7 deletions(-) create mode 100644 changelog.d/features/9415-newapi-sub2api-aggregator-balance.md create mode 100644 open-sse/services/newApiAggregatorQuotaFetcher.ts create mode 100644 src/app/(dashboard)/dashboard/providers/[id]/components/modals/NewApiAggregatorFields.tsx create mode 100644 tests/unit/newapi-aggregator-preflight-dispatch.test.ts create mode 100644 tests/unit/newapi-aggregator-quota-fetcher.test.ts diff --git a/changelog.d/features/9415-newapi-sub2api-aggregator-balance.md b/changelog.d/features/9415-newapi-sub2api-aggregator-balance.md new file mode 100644 index 0000000000..421c33b198 --- /dev/null +++ b/changelog.d/features/9415-newapi-sub2api-aggregator-balance.md @@ -0,0 +1,6 @@ +--- +kind: feature +ref: "#9415" +--- + +New-API / One-API / Sub2API aggregator balance detection for compatible nodes. When a compatible provider node has the "Aggregator Gateway" toggle enabled, OmniRoute will query the aggregator's `/api/user/self` endpoint to detect the account balance. The dashboard shows the balance badge and quota-preflight routing skips exhausted accounts. The feature is gated by the `NEWAPI_AGGREGATOR_BALANCE` feature flag (default: off). A custom `quotaPerUnit` override is supported for aggregators that use a different rate than the default 500000 units/$1. diff --git a/open-sse/services/newApiAggregatorQuotaFetcher.ts b/open-sse/services/newApiAggregatorQuotaFetcher.ts new file mode 100644 index 0000000000..cdcfca3274 --- /dev/null +++ b/open-sse/services/newApiAggregatorQuotaFetcher.ts @@ -0,0 +1,225 @@ +/** + * newApiAggregatorQuotaFetcher.ts — Generalized New-API / One-API / Sub2API + * Aggregator Balance Quota Fetcher + * + * Generalizes the AgentRouter (agentrouterQuotaFetcher.ts) balance detection + * so any OpenAI/Anthropic-compatible custom node pointing at a self-hosted + * New-API / One-API / Sub2API gateway can report its balance. + * + * New-API (QuantumNous/new-api, a fork of One API) exposes: + * + * GET {base}/api/user/self + * Authorization: Bearer {systemAccessToken} + * New-Api-User: {userId} + * -> { "data": { "quota": } } (raw New-API credit units) + * + * `quota_per_unit` (units per $1) defaults to 500000, overridable via + * `providerSpecificData.quotaPerUnit`. + * + * Credentials: the System Access Token + New-Api-User id are read from + * `connection.providerSpecificData.consoleApiKey` (reusing the existing generic + * field, same precedent as AgentRouter/Bailian) and + * `connection.providerSpecificData.newApiUserId` respectively. + * + * The `newApiAggregatorBalance` boolean flag in providerSpecificData must be + * `true` for the fetcher to activate — this is the opt-in toggle. + * + * Cache: in-memory TTL (60s), same pattern as sibling fetchers. + * + * Registration: this module exports fetchNewApiAggregatorQuota for dynamic + * dispatch; it does NOT self-register against a static provider key. + * Dynamic dispatch is handled by quotaPreflight.ts + quotaMonitor.ts. + */ + +import type { QuotaInfo } from "./quotaPreflight.ts"; +import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { toNumber } from "@/shared/utils/numeric"; + +const SELF_PATH = "/api/user/self"; + +// New-API-wide default: units per $1. See #6850 — can be hardcoded rather +// than fetched from /api/status on every call. +const DEFAULT_QUOTA_PER_UNIT = 500_000; + +const CACHE_TTL_MS = 60_000; // 60 seconds + +export interface NewApiAggregatorQuota extends QuotaInfo { + rawQuota: number; + dollarBalance: number; + limitReached: boolean; +} + +interface CacheEntry { + quota: NewApiAggregatorQuota; + fetchedAt: number; +} + +const quotaCache = new Map(); + +const _cacheCleanup = setInterval(() => { + const now = Date.now(); + for (const [key, entry] of quotaCache) { + if (now - entry.fetchedAt > CACHE_TTL_MS * 5) { + quotaCache.delete(key); + } + } +}, 5 * 60_000); + +if (typeof _cacheCleanup === "object" && "unref" in _cacheCleanup) { + (_cacheCleanup as { unref?: () => void }).unref?.(); +} + +function toRecord(value: unknown): Record { + return value && typeof value === "object" && !Array.isArray(value) + ? (value as Record) + : {}; +} + +/** + * Strip trailing `/v1` (or `/v1/`) from a baseUrl so that node URLs like + * `https://host/v1` still hit `{host}/api/user/self` rather than + * `{host}/v1/api/user/self`. + */ +function stripV1Suffix(baseUrl: string): string { + return baseUrl.replace(/\/v1\/?$/, ""); +} + +function extractCredentials(connection?: Record): { + systemAccessToken: string | null; + userId: string | null; + baseUrl: string | null; + quotaPerUnit: number; + aggregatorFlag: boolean; +} { + const providerSpecificData = toRecord(connection?.providerSpecificData); + const systemAccessToken = + typeof providerSpecificData.consoleApiKey === "string" && + providerSpecificData.consoleApiKey.trim().length > 0 + ? providerSpecificData.consoleApiKey + : null; + const userId = + typeof providerSpecificData.newApiUserId === "string" && + providerSpecificData.newApiUserId.trim().length > 0 + ? providerSpecificData.newApiUserId + : null; + const rawBaseUrl = + typeof providerSpecificData.baseUrl === "string" && + providerSpecificData.baseUrl.trim().length > 0 + ? providerSpecificData.baseUrl.trim() + : null; + const baseUrl = rawBaseUrl ? stripV1Suffix(rawBaseUrl) : null; + + const rawQuotaPerUnit = toNumber(providerSpecificData.quotaPerUnit, 0); + const quotaPerUnit = rawQuotaPerUnit > 0 ? rawQuotaPerUnit : DEFAULT_QUOTA_PER_UNIT; + + const aggregatorFlag = providerSpecificData.newApiAggregatorBalance === true; + + return { systemAccessToken, userId, baseUrl, quotaPerUnit, aggregatorFlag }; +} + +function parseNewApiAggregatorQuotaResponse( + data: unknown, + quotaPerUnit: number +): NewApiAggregatorQuota | null { + const obj = toRecord(data); + const dataObj = toRecord(obj.data); + + const rawQuotaValue = "quota" in dataObj ? dataObj.quota : obj.quota; + if (rawQuotaValue === undefined) return null; + + const rawQuota = toNumber(rawQuotaValue, -1); + if (rawQuota < 0) return null; + + const dollarBalance = rawQuota / quotaPerUnit; + const limitReached = rawQuota <= 0; + // No known upstream "total" grant to compute a real percentage against — follow + // DeepSeek's boolean-availability precedent (0% used = has balance, 100% = exhausted). + const percentUsed = limitReached ? 1 : 0; + + return { + used: percentUsed * 100, + total: 100, + percentUsed, + resetAt: null, + rawQuota, + dollarBalance, + limitReached, + }; +} + +/** + * Fetch current quota for a New-API / One-API / Sub2API aggregator connection. + * + * @param connectionId - Connection ID from the DB (used to key the cache) + * @param connection - Optional connection object with providerSpecificData credentials + * @returns NewApiAggregatorQuota or null if fetch fails / no credentials / not opted in + */ +export async function fetchNewApiAggregatorQuota( + connectionId: string, + connection?: Record +): Promise { + const cached = quotaCache.get(connectionId); + if (cached && Date.now() - cached.fetchedAt < CACHE_TTL_MS) { + return cached.quota; + } + + const { systemAccessToken, userId, baseUrl, quotaPerUnit, aggregatorFlag } = + extractCredentials(connection); + + if (!aggregatorFlag) return null; + if (!systemAccessToken || !userId || !baseUrl) return null; + + const url = `${baseUrl}${SELF_PATH}`; + + try { + await throttleQuotaFetch(); + + const response = await fetch(url, { + method: "GET", + headers: { + Authorization: `Bearer ${systemAccessToken}`, + "New-Api-User": userId, + "Content-Type": "application/json", + Accept: "application/json", + }, + signal: AbortSignal.timeout(8_000), + }); + + if (response.status === 401 || response.status === 403) { + quotaCache.delete(connectionId); + return null; + } + + if (!response.ok) { + return null; + } + + const data = await response.json(); + const quota = parseNewApiAggregatorQuotaResponse(data, quotaPerUnit); + + if (!quota) return null; + + quotaCache.set(connectionId, { quota, fetchedAt: Date.now() }); + return quota; + } catch { + return null; + } +} + +/** + * Force-invalidate the cache for a connection. + */ +export function invalidateNewApiAggregatorQuotaCache(connectionId: string): void { + quotaCache.delete(connectionId); +} + +/** + * Check whether a connection has opted in to New-API aggregator balance + * detection. Used by the dynamic dispatch in quotaPreflight / quotaMonitor. + */ +export function isNewApiAggregatorBalanceConnection( + connection?: Record +): boolean { + const providerSpecificData = toRecord(connection?.providerSpecificData); + return providerSpecificData.newApiAggregatorBalance === true; +} diff --git a/open-sse/services/quotaMonitor.ts b/open-sse/services/quotaMonitor.ts index da647c5f73..7fe91cd865 100644 --- a/open-sse/services/quotaMonitor.ts +++ b/open-sse/services/quotaMonitor.ts @@ -8,7 +8,11 @@ * Alertas deduplicados por sessão (janela de 5min). */ -import { registerQuotaFetcher, type QuotaFetcher } from "./quotaPreflight.ts"; +import { + registerQuotaFetcher, + resolveDynamicQuotaFetcher, + type QuotaFetcher, +} from "./quotaPreflight.ts"; import { getSessionInfo } from "./sessionManager.ts"; export { registerQuotaFetcher }; @@ -199,7 +203,12 @@ function scheduleNextPoll(sessionId: string, intervalMs: number): void { } try { - const fetcher = quotaFetcherRegistry.get(provider); + let fetcher = quotaFetcherRegistry.get(provider); + // Dynamic fallback: for compatible-provider connections with the + // aggregator flag + feature flag, use the generalized New-API fetcher. + if (!fetcher && current.connectionSnapshot) { + fetcher = resolveDynamicQuotaFetcher(provider, current.connectionSnapshot); + } if (!fetcher) { current.status = current.lastQuotaPercent === null ? "idle" : current.status; scheduleNextPoll(sessionId, NORMAL_INTERVAL_MS); diff --git a/open-sse/services/quotaPreflight.ts b/open-sse/services/quotaPreflight.ts index 9b0d7b8233..a7db736c29 100644 --- a/open-sse/services/quotaPreflight.ts +++ b/open-sse/services/quotaPreflight.ts @@ -18,6 +18,10 @@ * it — once you invoke preflight, it runs the fetcher and evaluates. */ +import { isCompatibleProviderConnectionId } from "@/shared/utils/compatibleProviderId"; +import { isFeatureFlagEnabled } from "@/shared/utils/featureFlags"; +import { fetchNewApiAggregatorQuota } from "./newApiAggregatorQuotaFetcher.ts"; + export interface PreflightQuotaResult { proceed: boolean; reason?: string; @@ -231,6 +235,29 @@ export function evaluateQuotaCutoff( return quotaPercentCutoffResult(quota, thresholds); } +/** + * Resolve a dynamic quota fetcher for compatible-provider connections that + * opt in to New-API / One-API / Sub2API aggregator balance detection. + * Returns the fetcher when both the feature flag and the connection's + * aggregator flag are true; otherwise returns undefined. + */ +export function resolveDynamicQuotaFetcher( + provider: string, + connection: Record +): QuotaFetcher | undefined { + // Dynamic dispatch only for compatible-provider connection IDs + if (!isCompatibleProviderConnectionId(provider)) return undefined; + + // Connection must opt in via providerSpecificData.newApiAggregatorBalance + const psd = connection?.providerSpecificData as Record | undefined; + if (!psd || psd.newApiAggregatorBalance !== true) return undefined; + + // Feature flag must be enabled + if (!isFeatureFlagEnabled("NEWAPI_AGGREGATOR_BALANCE")) return undefined; + + return fetchNewApiAggregatorQuota; +} + export async function preflightQuota( provider: string, connectionId: string, @@ -239,9 +266,14 @@ export async function preflightQuota( ): Promise { // No legacy enable-flag gate here — the caller decides when to invoke us // (see file-level docstring). When there's no fetcher we proceed silently. - const fetcher = getQuotaFetcher(provider); + let fetcher = getQuotaFetcher(provider); if (!fetcher) { - return { proceed: true }; + // Dynamic fallback: for compatible-provider connections with the + // aggregator flag + feature flag, use the generalized New-API fetcher. + fetcher = resolveDynamicQuotaFetcher(provider, connection); + if (!fetcher) { + return { proceed: true }; + } } let quota: QuotaInfo | null = null; diff --git a/src/app/(dashboard)/dashboard/providers/[id]/components/modals/AddApiKeyModal.tsx b/src/app/(dashboard)/dashboard/providers/[id]/components/modals/AddApiKeyModal.tsx index cc8e7df52c..80ca8af473 100644 --- a/src/app/(dashboard)/dashboard/providers/[id]/components/modals/AddApiKeyModal.tsx +++ b/src/app/(dashboard)/dashboard/providers/[id]/components/modals/AddApiKeyModal.tsx @@ -123,6 +123,8 @@ export default function AddApiKeyModal({ accountId: "", consoleApiKey: "", newApiUserId: "", + newApiAggregatorBalance: false, + quotaPerUnit: "", ...EMPTY_GLM_TEAM_QUOTA_FIELDS, ...EMPTY_QUOTA_SCRAPING_FIELDS, ccCompatibleContext1m: false, diff --git a/src/app/(dashboard)/dashboard/providers/[id]/components/modals/EditCompatibleNodeModal.tsx b/src/app/(dashboard)/dashboard/providers/[id]/components/modals/EditCompatibleNodeModal.tsx index 2d01e82b71..2c545471e1 100644 --- a/src/app/(dashboard)/dashboard/providers/[id]/components/modals/EditCompatibleNodeModal.tsx +++ b/src/app/(dashboard)/dashboard/providers/[id]/components/modals/EditCompatibleNodeModal.tsx @@ -1,8 +1,9 @@ "use client"; import { useState, useEffect } from "react"; import { useTranslations } from "next-intl"; -import { Button, Badge, Input, Modal, Select } from "@/shared/components"; +import { Button, Badge, Input, Modal, Select, Toggle } from "@/shared/components"; import { CC_COMPATIBLE_DEFAULT_CHAT_PATH } from "../../providerDetailConstants"; +import NewApiAggregatorFields from "./NewApiAggregatorFields"; interface EditCompatibleNodeModalNode { id?: string; name?: string; @@ -12,6 +13,7 @@ interface EditCompatibleNodeModalNode { chatPath?: string; modelsPath?: string; iconUrl?: string; + providerSpecificData?: Record; } interface EditCompatibleNodeModalProps { @@ -40,6 +42,10 @@ export default function EditCompatibleNodeModal({ chatPath: "", modelsPath: "", iconUrl: "", + newApiAggregatorBalance: false, + consoleApiKey: "", + newApiUserId: "", + quotaPerUnit: "", }); const [saving, setSaving] = useState(false); const [checkKey, setCheckKey] = useState(""); @@ -54,6 +60,7 @@ export default function EditCompatibleNodeModal({ useEffect(() => { if (node) { + const psd = (node.providerSpecificData || {}) as Record; setFormData({ name: node.name || "", prefix: node.prefix || "", @@ -68,6 +75,10 @@ export default function EditCompatibleNodeModal({ chatPath: node.chatPath || (isCcCompatible ? CC_COMPATIBLE_DEFAULT_CHAT_PATH : ""), modelsPath: isCcCompatible ? "" : node.modelsPath || "", iconUrl: node.iconUrl || "", + newApiAggregatorBalance: psd.newApiAggregatorBalance === true, + consoleApiKey: typeof psd.consoleApiKey === "string" ? psd.consoleApiKey : "", + newApiUserId: typeof psd.newApiUserId === "string" ? psd.newApiUserId : "", + quotaPerUnit: typeof psd.quotaPerUnit === "number" ? String(psd.quotaPerUnit) : "", }); setShowAdvanced( !!( @@ -103,6 +114,22 @@ export default function EditCompatibleNodeModal({ if (!isAnthropic) { payload.apiType = formData.apiType; } + // Aggregator gateway fields (#9415) + if (formData.newApiAggregatorBalance) { + payload.providerSpecificData = { + newApiAggregatorBalance: true, + }; + if (formData.consoleApiKey.trim()) { + payload.providerSpecificData.consoleApiKey = formData.consoleApiKey.trim(); + } + if (formData.newApiUserId.trim()) { + payload.providerSpecificData.newApiUserId = formData.newApiUserId.trim(); + } + const parsedQuotaPerUnit = parseInt(formData.quotaPerUnit, 10); + if (Number.isFinite(parsedQuotaPerUnit) && parsedQuotaPerUnit > 0) { + payload.providerSpecificData.quotaPerUnit = parsedQuotaPerUnit; + } + } await onSave(payload); } finally { setSaving(false); @@ -222,6 +249,24 @@ export default function EditCompatibleNodeModal({ placeholder="https://example.com/logo.png" hint={t("iconUrlHint")} /> + + setFormData({ ...formData, newApiAggregatorBalance: checked }) + } + /> + setFormData({ ...formData, ...patch })} + t={t} + />