From d31ab964504ea1fab71eee94ebb9970c8f77106b Mon Sep 17 00:00:00 2001 From: Dizzle <112548150+maxmad64bis@users.noreply.github.com> Date: Sun, 11 Oct 2026 14:19:13 +0200 Subject: [PATCH 1/3] fix(quota): route quota fetcher reads through the connection proxy Console quota reads used a direct fetch; they now run inside the assigned connection proxy context with no direct fallback. --- ...quota-fetchers-through-connection-proxy.md | 1 + open-sse/services/agentrouterQuotaFetcher.ts | 4 +- .../services/alibabaFreeTierQuotaFetcher.ts | 53 +- open-sse/services/bailianQuotaFetcher.ts | 7 +- open-sse/services/codexQuotaFetcher.ts | 4 +- open-sse/services/context7QuotaFetcher.ts | 4 +- open-sse/services/deepseekQuotaFetcher.ts | 4 +- open-sse/services/firecrawlQuotaFetcher.ts | 4 +- open-sse/services/grokCliQuotaFetcher.ts | 4 +- open-sse/services/grokQuotaFetcher.ts | 54 +- .../services/qwenTokenPlanQuotaFetcher.ts | 14 +- open-sse/services/xaiOauthQuotaFetcher.ts | 2 +- tests/unit/bailian-usage.test.ts | 23 + .../quota-fetchers-connection-proxy.test.ts | 580 ++++++++++++++++++ 14 files changed, 709 insertions(+), 49 deletions(-) create mode 100644 changelog.d/fixes/15885-quota-fetchers-through-connection-proxy.md create mode 100644 tests/unit/quota-fetchers-connection-proxy.test.ts diff --git a/changelog.d/fixes/15885-quota-fetchers-through-connection-proxy.md b/changelog.d/fixes/15885-quota-fetchers-through-connection-proxy.md new file mode 100644 index 000000000000..ea269e6da9b0 --- /dev/null +++ b/changelog.d/fixes/15885-quota-fetchers-through-connection-proxy.md @@ -0,0 +1 @@ +- **fix(quota):** quota readings use the assigned connection egress path ([#15885](https://github.com/diegosouzapw/OmniRoute/pull/15885)) — thanks @maxmad64bis diff --git a/open-sse/services/agentrouterQuotaFetcher.ts b/open-sse/services/agentrouterQuotaFetcher.ts index 58672b676c14..9869fa21bfc2 100644 --- a/open-sse/services/agentrouterQuotaFetcher.ts +++ b/open-sse/services/agentrouterQuotaFetcher.ts @@ -28,6 +28,7 @@ import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; const AGENTROUTER_CONFIG = { baseUrl: "https://agentrouter.org", @@ -152,7 +153,7 @@ export async function fetchAgentrouterQuota( try { await throttleQuotaFetch(); - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers: { Authorization: `Bearer ${systemAccessToken}`, @@ -162,6 +163,7 @@ export async function fetchAgentrouterQuota( }, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; if (response.status === 401 || response.status === 403) { quotaCache.delete(connectionId); diff --git a/open-sse/services/alibabaFreeTierQuotaFetcher.ts b/open-sse/services/alibabaFreeTierQuotaFetcher.ts index 4cf39a6e4320..224a9cc75971 100644 --- a/open-sse/services/alibabaFreeTierQuotaFetcher.ts +++ b/open-sse/services/alibabaFreeTierQuotaFetcher.ts @@ -27,6 +27,7 @@ */ import { getAlibabaBillingMode, isAlibabaModelStudioProvider } from "./alibabaFreeTier.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { asRecord, getAlibabaFreeTierQuotaLastSyncAt, @@ -302,6 +303,7 @@ function buildRequestBody( } async function postConsoleFreeTierQuota( + connectionId: string | null, region: AlibabaProviderRegion, api: string, cookie: string, @@ -314,7 +316,10 @@ async function postConsoleFreeTierQuota( body.set("sec_token", secToken); } - const response = await fetch(buildGatewayUrl(region, api), { + const response = await fetchWithConnectionProxy( + connectionId ?? "", + buildGatewayUrl(region, api), + { method: "POST", headers: { Accept: "*/*", @@ -328,6 +333,7 @@ async function postConsoleFreeTierQuota( body: body.toString(), signal: AbortSignal.timeout(15_000), }); + if (!response) return null; return response.json(); } @@ -350,7 +356,8 @@ async function delay(ms: number): Promise { async function fetchAlibabaFreeTierQuotaEntriesForPath( providerSpecificData: Record | null | undefined, - fePath: string + fePath: string, + connectionId?: string | null ): Promise { const cookie = getAlibabaConsoleCookie(providerSpecificData); if (!cookie) return null; @@ -359,6 +366,7 @@ async function fetchAlibabaFreeTierQuotaEntriesForPath( const secToken = getAlibabaConsoleSecToken(providerSpecificData); let payload = await postConsoleFreeTierQuota( + connectionId ?? null, region, FREE_TIER_QUOTA_START_API, cookie, @@ -369,6 +377,7 @@ async function fetchAlibabaFreeTierQuotaEntriesForPath( if (!hasQuotaPayload(payload)) { payload = await postConsoleFreeTierQuota( + connectionId ?? null, region, FREE_TIER_QUOTA_API, cookie, @@ -387,6 +396,7 @@ async function fetchAlibabaFreeTierQuotaEntriesForPath( await delay(400); } payload = await postConsoleFreeTierQuota( + connectionId ?? null, region, FREE_TIER_QUOTA_API, cookie, @@ -403,27 +413,31 @@ async function fetchAlibabaFreeTierQuotaEntriesForPath( } export async function fetchAlibabaFreeTierQuotaEntries( - providerSpecificData: Record | null | undefined + providerSpecificData: Record | null | undefined, + connectionId?: string | null ): Promise { - return fetchAlibabaFreeTierQuotaEntriesForPath(providerSpecificData, DEFAULT_TEXT_FE_PATH); + return fetchAlibabaFreeTierQuotaEntriesForPath(providerSpecificData, DEFAULT_TEXT_FE_PATH, connectionId); } export async function fetchAlibabaFreeTierVisionQuotaEntries( - providerSpecificData: Record | null | undefined + providerSpecificData: Record | null | undefined, + connectionId?: string | null ): Promise { - return fetchAlibabaFreeTierQuotaEntriesForPath(providerSpecificData, DEFAULT_VISION_FE_PATH); + return fetchAlibabaFreeTierQuotaEntriesForPath(providerSpecificData, DEFAULT_VISION_FE_PATH, connectionId); } export async function fetchAlibabaFreeTierMultimodalQuotaEntries( - providerSpecificData: Record | null | undefined + providerSpecificData: Record | null | undefined, + connectionId?: string | null ): Promise { - return fetchAlibabaFreeTierQuotaEntriesForPath(providerSpecificData, DEFAULT_MULTIMODAL_FE_PATH); + return fetchAlibabaFreeTierQuotaEntriesForPath(providerSpecificData, DEFAULT_MULTIMODAL_FE_PATH, connectionId); } export async function fetchAlibabaFreeTierAudioQuotaEntries( - providerSpecificData: Record | null | undefined + providerSpecificData: Record | null | undefined, + connectionId?: string | null ): Promise { - return fetchAlibabaFreeTierQuotaEntriesForPath(providerSpecificData, DEFAULT_AUDIO_FE_PATH); + return fetchAlibabaFreeTierQuotaEntriesForPath(providerSpecificData, DEFAULT_AUDIO_FE_PATH, connectionId); } function mergeUniqueQuotaEntries( @@ -442,17 +456,18 @@ function mergeUniqueQuotaEntries( } export async function buildAlibabaFreeTierQuotaSnapshot( - providerSpecificData: Record | null | undefined + providerSpecificData: Record | null | undefined, + connectionId?: string | null ): Promise { - const textEntries = await fetchAlibabaFreeTierQuotaEntries(providerSpecificData); + const textEntries = await fetchAlibabaFreeTierQuotaEntries(providerSpecificData, connectionId); if (!textEntries) return null; const visionEntries = - (await fetchAlibabaFreeTierVisionQuotaEntries(providerSpecificData)) || textEntries; + (await fetchAlibabaFreeTierVisionQuotaEntries(providerSpecificData, connectionId)) || textEntries; const multimodalEntries = - (await fetchAlibabaFreeTierMultimodalQuotaEntries(providerSpecificData)) || textEntries; + (await fetchAlibabaFreeTierMultimodalQuotaEntries(providerSpecificData, connectionId)) || textEntries; const audioEntries = - (await fetchAlibabaFreeTierAudioQuotaEntries(providerSpecificData)) || textEntries; + (await fetchAlibabaFreeTierAudioQuotaEntries(providerSpecificData, connectionId)) || textEntries; const text = classifyAlibabaFreeTierQuotaEntries(textEntries, { textOnly: true }); const vision = classifyAlibabaVisionFreeTierQuotaEntries(visionEntries); @@ -470,7 +485,8 @@ export async function buildAlibabaFreeTierQuotaSnapshot( export async function refreshAlibabaFreeTierQuotaClassification( provider: string, - providerSpecificData: Record | null | undefined + providerSpecificData: Record | null | undefined, + connectionId?: string | null ): Promise | null> { if ( !isAlibabaModelStudioProvider(provider) || @@ -482,7 +498,7 @@ export async function refreshAlibabaFreeTierQuotaClassification( return null; } - const snapshot = await buildAlibabaFreeTierQuotaSnapshot(providerSpecificData); + const snapshot = await buildAlibabaFreeTierQuotaSnapshot(providerSpecificData, connectionId); if (!snapshot) return null; return mergeAlibabaFreeTierQuotaClassification(providerSpecificData, snapshot); @@ -503,7 +519,8 @@ export function scheduleAlibabaFreeTierQuotaRefresh( try { const merged = await refreshAlibabaFreeTierQuotaClassification( provider, - connection.providerSpecificData + connection.providerSpecificData, + connection.id ); if (!merged) return; const { updateProviderConnection } = await import("../../src/lib/db/providers.ts"); diff --git a/open-sse/services/bailianQuotaFetcher.ts b/open-sse/services/bailianQuotaFetcher.ts index 921bd3472177..c0771d2acb68 100644 --- a/open-sse/services/bailianQuotaFetcher.ts +++ b/open-sse/services/bailianQuotaFetcher.ts @@ -22,6 +22,7 @@ import { registerQuotaFetcher, registerQuotaWindows, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; // Bailian quota hosts (international / china fallback) const BAILIAN_QUOTA_HOSTS = { @@ -259,13 +260,14 @@ export async function fetchBailianQuota( const url = getQuotaUrl(); // #6911: space concurrent upstream quota fetches (mirrors codexQuotaFetcher.ts). await throttleQuotaFetch(); - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "POST", headers, body: JSON.stringify({}), signal: AbortSignal.timeout(8_000), }); + if (!response) return null; const rawData = await response.json(); const obj = toRecord(rawData); @@ -278,12 +280,13 @@ export async function fetchBailianQuota( // #6911: space this fallback fetch too — it is still a genuine upstream call. await throttleQuotaFetch(); - const retryResponse = await fetch(chinaUrl, { + const retryResponse = await fetchWithConnectionProxy(connectionId, chinaUrl, { method: "POST", headers, body: JSON.stringify({}), signal: AbortSignal.timeout(8_000), }); + if (!retryResponse) return null; const retryData = await retryResponse.json(); const quota = parseBailianQuotaResponse(retryData); diff --git a/open-sse/services/codexQuotaFetcher.ts b/open-sse/services/codexQuotaFetcher.ts index a0b0bcf91bb1..00ffc11db1b1 100644 --- a/open-sse/services/codexQuotaFetcher.ts +++ b/open-sse/services/codexQuotaFetcher.ts @@ -25,6 +25,7 @@ import { import { registerQuotaFetcher, registerQuotaWindows, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { getCodexBackendIdentityHeaders } from "../config/codexClient.ts"; import { parseCodexPaidCredits } from "@/lib/providers/codexPaidCredits"; @@ -240,11 +241,12 @@ export async function fetchCodexQuota( // configurable (OMNIROUTE_QUOTA_FETCH_MIN_INTERVAL_MS, 0 = disabled). await throttleQuotaFetch(); - const response = await fetch(CODEX_USAGE_URL, { + const response = await fetchWithConnectionProxy(connectionId, CODEX_USAGE_URL, { method: "GET", headers, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; if (!response.ok) { // Non-2xx: could be token expired or quota API down. diff --git a/open-sse/services/context7QuotaFetcher.ts b/open-sse/services/context7QuotaFetcher.ts index d24bbfd9966c..a6ac9213bbb6 100644 --- a/open-sse/services/context7QuotaFetcher.ts +++ b/open-sse/services/context7QuotaFetcher.ts @@ -1,6 +1,7 @@ import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; /** * Context7 has no dedicated usage/billing endpoint, so quota is read from the @@ -139,11 +140,12 @@ export async function fetchContext7Quota( return dedupeQuotaFetch(connectionId, async () => { await throttleQuotaFetch(); const url = `${CONTEXT7_SEARCH_URL}?query=react`; - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers: { Authorization: `Bearer ${token}` }, signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), }); + if (!response) return null; const quota = parseContext7RateLimitHeaders(response.headers); if (quota) { diff --git a/open-sse/services/deepseekQuotaFetcher.ts b/open-sse/services/deepseekQuotaFetcher.ts index f7f771f58c47..5d20575e5d9c 100644 --- a/open-sse/services/deepseekQuotaFetcher.ts +++ b/open-sse/services/deepseekQuotaFetcher.ts @@ -25,6 +25,7 @@ import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; // DeepSeek API config const DEEPSEEK_CONFIG = { @@ -194,7 +195,7 @@ export async function fetchDeepseekQuota( // Cache hits above never reach here; this only paces genuine network calls. await throttleQuotaFetch(); - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers: { Authorization: `Bearer ${apiKey}`, @@ -203,6 +204,7 @@ export async function fetchDeepseekQuota( }, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; // 401/403: token invalid — remove from cache if (response.status === 401 || response.status === 403) { diff --git a/open-sse/services/firecrawlQuotaFetcher.ts b/open-sse/services/firecrawlQuotaFetcher.ts index f53229051131..dbb6ea04418b 100644 --- a/open-sse/services/firecrawlQuotaFetcher.ts +++ b/open-sse/services/firecrawlQuotaFetcher.ts @@ -25,6 +25,7 @@ import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { isUrlHost } from "../utils/urlHost.ts"; import { toNumberOrNull } from "@/shared/utils/numeric"; @@ -156,7 +157,7 @@ export async function fetchFirecrawlQuota( try { await throttleQuotaFetch(); - const response = await fetch(CREDIT_USAGE_URL, { + const response = await fetchWithConnectionProxy(connectionId, CREDIT_USAGE_URL, { method: "GET", headers: { Authorization: `Bearer ${apiKey}`, @@ -164,6 +165,7 @@ export async function fetchFirecrawlQuota( }, signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), }); + if (!response) return null; if (response.status === 401 || response.status === 403) { quotaCache.set(connectionId, { quota: null, fetchedAt: Date.now() }); diff --git a/open-sse/services/grokCliQuotaFetcher.ts b/open-sse/services/grokCliQuotaFetcher.ts index e24b97191f15..16157aca0eb9 100644 --- a/open-sse/services/grokCliQuotaFetcher.ts +++ b/open-sse/services/grokCliQuotaFetcher.ts @@ -36,6 +36,7 @@ import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { decodeGrokCreditsFrame } from "./grokCliQuotaFrame.ts"; const GROK_CLI_CONFIG = { @@ -142,7 +143,7 @@ export async function fetchGrokCliQuota( try { await throttleQuotaFetch(); - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "POST", headers: { Authorization: `Bearer ${accessToken}`, @@ -152,6 +153,7 @@ export async function fetchGrokCliQuota( body: GRPC_WEB_EMPTY_REQUEST_FRAME, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; if (response.status === 401 || response.status === 403) { // Fail-open: no proactive refreshCredentials() in this PR (see Non-Goals). diff --git a/open-sse/services/grokQuotaFetcher.ts b/open-sse/services/grokQuotaFetcher.ts index 185c7e99a6bf..6cfaddf44165 100644 --- a/open-sse/services/grokQuotaFetcher.ts +++ b/open-sse/services/grokQuotaFetcher.ts @@ -21,6 +21,7 @@ */ import { registerQuotaFetcher, registerQuotaWindows, type QuotaInfo } from "./quotaPreflight.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { existsSync, readFileSync, writeFileSync } from "node:fs"; import { homedir } from "node:os"; @@ -204,15 +205,20 @@ function writeRefreshedTokens( } } -async function discoverTokenEndpoint(issuer: string, signal: AbortSignal): Promise { +async function discoverTokenEndpoint( + issuer: string, + signal: AbortSignal, + connectionId?: string | null +): Promise { const discoveryUrl = `${issuer.replace(/\/$/, "")}/.well-known/openid-configuration`; if (!isAllowedXaiUrl(discoveryUrl)) { throw new Error("invalid oidc issuer"); } - const res = await fetch(discoveryUrl, { + const res = await fetchWithConnectionProxy(connectionId ?? "", discoveryUrl, { headers: { Accept: "application/json" }, signal, }); + if (!res) throw new Error("proxy quota fetch failed"); if (!res.ok) throw new Error(`token refresh failed (discovery HTTP ${res.status})`); const json = (await res.json()) as { token_endpoint?: string }; const endpoint = String(json.token_endpoint || ""); @@ -220,7 +226,11 @@ async function discoverTokenEndpoint(issuer: string, signal: AbortSignal): Promi return endpoint; } -async function refreshAccessToken(auth: ResolvedAuth, signal: AbortSignal): Promise { +async function refreshAccessToken( + auth: ResolvedAuth, + signal: AbortSignal, + connectionId?: string | null +): Promise { if (!auth.refreshToken) { throw new Error("auth expired — run `grok login`"); } @@ -228,8 +238,8 @@ async function refreshAccessToken(auth: ResolvedAuth, signal: AbortSignal): Prom throw new Error("auth missing client id — run `grok login`"); } - const tokenEndpoint = await discoverTokenEndpoint(auth.issuer, signal); - const res = await fetch(tokenEndpoint, { + const tokenEndpoint = await discoverTokenEndpoint(auth.issuer, signal, connectionId); + const res = await fetchWithConnectionProxy(connectionId ?? "", tokenEndpoint, { method: "POST", headers: { "Content-Type": "application/x-www-form-urlencoded", @@ -238,11 +248,12 @@ async function refreshAccessToken(auth: ResolvedAuth, signal: AbortSignal): Prom }, body: new URLSearchParams({ grant_type: "refresh_token", - client_id: auth.clientId, - refresh_token: auth.refreshToken, + client_id: auth.clientId ?? "", + refresh_token: auth.refreshToken ?? "", }).toString(), signal, }); + if (!res) throw new Error("proxy quota fetch failed"); if (!res.ok) { throw new Error(`token refresh failed (HTTP ${res.status})`); @@ -286,13 +297,16 @@ function needsRefresh(auth: ResolvedAuth): boolean { return auth.expiresAtMs <= Date.now() + EXPIRY_SKEW_MS; } -async function resolveAuth(signal: AbortSignal): Promise { +async function resolveAuth( + signal: AbortSignal, + connectionId?: string | null +): Promise { const file = readAuthFile(); if (!file) throw new Error("no grok auth — run `grok login`"); const auth = pickAuthEntry(file); if (!auth) throw new Error("no usable Grok credentials — run `grok login`"); if (needsRefresh(auth)) { - return refreshAccessToken(auth, signal); + return refreshAccessToken(auth, signal, connectionId); } return auth; } @@ -329,12 +343,13 @@ function resetLocalLabel(endIso?: string): string { */ export async function fetchGrokBillingWithToken( token: string, - signal?: AbortSignal + signal?: AbortSignal, + connectionId?: string | null ): Promise { const controller = signal ? null : new AbortController(); const timeout = controller ? setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS) : null; try { - return await fetchBilling(token, signal ?? controller!.signal); + return await fetchBilling(token, signal ?? controller!.signal, connectionId); } finally { if (timeout) clearTimeout(timeout); } @@ -348,8 +363,12 @@ export function grokBillingSnapshotToQuotaInfo(snap: GrokBillingSnapshot): Quota return snapshotToQuotaInfo(snap); } -async function fetchBilling(token: string, signal: AbortSignal): Promise { - const res = await fetch(BILLING_URL, { +async function fetchBilling( + token: string, + signal: AbortSignal, + connectionId?: string | null +): Promise { + const res = await fetchWithConnectionProxy(connectionId ?? "", BILLING_URL, { method: "GET", headers: { Authorization: `Bearer ${token}`, @@ -359,6 +378,7 @@ async function fetchBilling(token: string, signal: AbortSignal): Promise controller.abort(), REQUEST_TIMEOUT_MS); try { - const snap = await fetchGrokBillingWithToken(accessToken, controller.signal); + const snap = await fetchGrokBillingWithToken(accessToken, controller.signal, connectionId); const quota = grokBillingSnapshotToQuotaInfo(snap); quotaCache.set(connectionId, { quota, fetchedAt: Date.now() }); return quota; diff --git a/tests/unit/bailian-usage.test.ts b/tests/unit/bailian-usage.test.ts index 20d7d4840848..b3abefaf4107 100644 --- a/tests/unit/bailian-usage.test.ts +++ b/tests/unit/bailian-usage.test.ts @@ -41,6 +41,7 @@ test("getUsageForProvider with bailian-coding-plan and consoleApiKey returns quo }; const result = await getUsageForProvider({ + id: "test-conn-xyz", provider: "bailian-coding-plan", apiKey: "sk-test", providerSpecificData: { consoleApiKey: "ck-test" }, @@ -94,6 +95,7 @@ test("getUsageForProvider with bailian-coding-plan and only apiKey falls back to }; const result = await getUsageForProvider({ + id: "test-conn-xyz", provider: "bailian-coding-plan", apiKey: "sk-test-fallback", providerSpecificData: {}, @@ -150,6 +152,7 @@ test("getUsageForProvider with bailian-coding-plan returns quota with percentUse }; const result = await getUsageForProvider({ + id: "test-conn-xyz", provider: "bailian-coding-plan", apiKey: "sk-test", providerSpecificData: { consoleApiKey: "ck-test" }, @@ -171,3 +174,23 @@ test("getUsageForProvider with bailian-coding-plan returns quota with percentUse "percentUsed should be 0.8 from most restrictive window (weekly 80%)" ); }); + +test("getUsageForProvider without a connection id reports no reading instead of fetching directly", async () => { + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return new Response(JSON.stringify({ code: "Success", data: {} }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }; + + const result = await getUsageForProvider({ + provider: "bailian-coding-plan", + apiKey: "sk-test", + providerSpecificData: { consoleApiKey: "ck-test" }, + }); + + assert.equal(result?.used, undefined, "Should not return quota without a connection id"); + assert.equal(fetchCalls, 0, "Should not fetch outside the connection proxy"); +}); diff --git a/tests/unit/quota-fetchers-connection-proxy.test.ts b/tests/unit/quota-fetchers-connection-proxy.test.ts new file mode 100644 index 000000000000..85c30d2b3de5 --- /dev/null +++ b/tests/unit/quota-fetchers-connection-proxy.test.ts @@ -0,0 +1,580 @@ +/** + * quota-fetchers-connection-proxy.test.ts — quota fetchers read through the + * connection proxy context. + * + * Behavioural guard: with a proxy assigned to the connection, the upstream + * quota request runs inside the connection proxy context instead of a direct + * fetch; failures stay visible (null) and never fall back to a direct fetch. + */ +import test from "node:test"; +import assert from "node:assert/strict"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +import { + connectionProxyTestOverride, + fetchWithConnectionProxy, + resolveConnectionProxy, +} from "../../open-sse/services/connectionProxyFetch.ts"; +import { hasAmbientProxyContext } from "../../open-sse/utils/proxyFetch.ts"; +import { + fetchAgentrouterQuota, + invalidateAgentrouterQuotaCache, +} from "../../open-sse/services/agentrouterQuotaFetcher.ts"; +import { + fetchAlibabaFreeTierQuotaEntries, + getAlibabaConsoleCookie, +} from "../../open-sse/services/alibabaFreeTierQuotaFetcher.ts"; +import { + fetchBailianQuota, + invalidateBailianQuotaCache, +} from "../../open-sse/services/bailianQuotaFetcher.ts"; +import { + fetchCodexQuota, + invalidateCodexQuotaCache, + registerCodexConnection, + unregisterCodexConnection, +} from "../../open-sse/services/codexQuotaFetcher.ts"; +import { + fetchContext7Quota, + invalidateContext7QuotaCache, +} from "../../open-sse/services/context7QuotaFetcher.ts"; +import { + fetchDeepseekQuota, + invalidateDeepseekQuotaCache, +} from "../../open-sse/services/deepseekQuotaFetcher.ts"; +import { + fetchFirecrawlQuota, + invalidateFirecrawlQuotaCache, +} from "../../open-sse/services/firecrawlQuotaFetcher.ts"; +import { + fetchGrokCliQuota, + invalidateGrokCliQuotaCache, +} from "../../open-sse/services/grokCliQuotaFetcher.ts"; +import { + fetchGrokWebQuota, + invalidateGrokWebQuotaCache, +} from "../../open-sse/services/grokQuotaFetcher.ts"; +import { + fetchQwenTokenPlanQuota, + invalidateQwenTokenPlanQuotaCache, +} from "../../open-sse/services/qwenTokenPlanQuotaFetcher.ts"; + +// socks5 skips the TCP reachability probe in runWithProxyContext, so the mock +// fetch below never races the fast-fail path; the proxy context is still set. +const ASSIGNED_PROXY = { type: "socks5", host: "127.0.0.1", port: 9 }; +const QUOTA_URL = "https://console.example/api/quota"; +const originalFetch = globalThis.fetch; +const originalAuthPath = process.env.GROK_AUTH_PATH; + +function quotaResponse(): Response { + return new Response(JSON.stringify({ quota: 1 }), { status: 200 }); +} + +/** Route every helper resolution through the assigned proxy in fetcher cases. */ +function useAssignedProxy(): void { + connectionProxyTestOverride.resolveProxyForConnection = async () => ({ + proxy: ASSIGNED_PROXY, + }); +} + +test.afterEach(() => { + globalThis.fetch = originalFetch; + connectionProxyTestOverride.resolveProxyForConnection = null; + if (originalAuthPath !== undefined) { + process.env.GROK_AUTH_PATH = originalAuthPath; + } else { + delete process.env.GROK_AUTH_PATH; + } + delete process.env.QWEN_CLOUD_COOKIE; + delete process.env.QWEN_CLOUD_SEC_TOKEN; +}); + +test("routes the quota fetch through the assigned connection proxy", async () => { + let observedInContext: boolean | null = null; + const fetchFn = async (): Promise => { + observedInContext = hasAmbientProxyContext(); + return quotaResponse(); + }; + const response = await fetchWithConnectionProxy("conn-proxy", QUOTA_URL, undefined, { + resolveProxyForConnection: async () => ({ proxy: ASSIGNED_PROXY }), + fetchFn: fetchFn as typeof globalThis.fetch, + }); + assert.ok(response instanceof Response); + assert.equal(observedInContext, true); +}); + +test("runs directly when the connection has no proxy assigned", async () => { + let observedInContext: boolean | null = null; + const fetchFn = async (): Promise => { + observedInContext = hasAmbientProxyContext(); + return quotaResponse(); + }; + const response = await fetchWithConnectionProxy("conn-direct", QUOTA_URL, undefined, { + resolveProxyForConnection: async () => ({ proxy: null }), + fetchFn: fetchFn as typeof globalThis.fetch, + }); + assert.ok(response instanceof Response); + assert.equal(observedInContext, false); +}); + +test("reports a missing connection identifier instead of fetching directly", async () => { + let calls = 0; + const fetchFn = async (): Promise => { + calls += 1; + return quotaResponse(); + }; + const response = await fetchWithConnectionProxy("", QUOTA_URL, undefined, { + resolveProxyForConnection: async () => ({ proxy: ASSIGNED_PROXY }), + fetchFn: fetchFn as typeof globalThis.fetch, + }); + assert.equal(response, null); + assert.equal(calls, 0); +}); + +test("reports a proxy failure instead of retrying outside the proxy", async () => { + let calls = 0; + const fetchFn = async (): Promise => { + calls += 1; + throw new Error("proxy blocked the console domain"); + }; + const response = await fetchWithConnectionProxy("conn-blocked", QUOTA_URL, undefined, { + resolveProxyForConnection: async () => ({ proxy: ASSIGNED_PROXY }), + fetchFn: fetchFn as typeof globalThis.fetch, + }); + assert.equal(response, null); + assert.equal(calls, 1); +}); + +test("reports a proxy resolution failure instead of fetching directly", async () => { + let calls = 0; + const fetchFn = async (): Promise => { + calls += 1; + return quotaResponse(); + }; + const response = await fetchWithConnectionProxy("conn-unknown", QUOTA_URL, undefined, { + resolveProxyForConnection: async () => { + throw new Error("unknown connection"); + }, + fetchFn: fetchFn as typeof globalThis.fetch, + }); + assert.equal(response, null); + assert.equal(calls, 0); +}); + +test("resolveConnectionProxy resolves through the stored connection binding", async () => { + const proxy = await resolveConnectionProxy("conn-proxy", { + resolveProxyForConnection: async () => ({ proxy: ASSIGNED_PROXY }), + }); + assert.deepEqual(proxy, ASSIGNED_PROXY); +}); + +test("resolveConnectionProxy returns null when no proxy is assigned", async () => { + const proxy = await resolveConnectionProxy("conn-direct", { + resolveProxyForConnection: async () => ({ proxy: null }), + }); + assert.equal(proxy, null); +}); + +function gatewayBody(payload: unknown, api: string): string { + return JSON.stringify({ + code: "200", + data: { + DataV2: { + ret: ["SUCCESS::ok"], + data: { msg: "Success.", code: "SUCCESS", data: payload, success: true }, + }, + success: true, + httpStatus: 200, + errorCode: "", + api, + errorMsg: "", + }, + httpStatusCode: "200", + successResponse: true, + }); +} + +test("qwen token plan quota reads the dashboard and gateway inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `qwen-proxy-${Date.now()}`; + const contexts: boolean[] = []; + const dashboardHtml = 'SEC_TOKEN: "sec-abc"'; + const usagePayload = { per1WeekResetTime: 1786714740000, per1WeekPercentage: 0.55 }; + const quotaConfig = { pro: { five_hour: 12000.0, weekly: 40000.0 } }; + const subscription = { specCode: "pro" }; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async (input: RequestInfo | URL) => { + contexts.push(hasAmbientProxyContext()); + const url = String(input); + if (!url.includes("/data/api.json")) { + return new Response(dashboardHtml, { status: 200 }); + } + if (url.includes("%2Fusage")) { + return new Response(gatewayBody(usagePayload, "usage"), { status: 200 }); + } + if (url.includes("%2Fquota-config")) { + return new Response(gatewayBody(quotaConfig, "quota-config"), { status: 200 }); + } + return new Response(gatewayBody(subscription, "subscription"), { status: 200 }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchQwenTokenPlanQuota(connectionId, { + providerSpecificData: { qwenCloudCookie: "login_qwencloud_ticket=tick" }, + provider: "qwen-cloud-token-plan", + }); + assert.ok(quota); + assert.equal(contexts.length >= 2, true); + assert.ok(contexts.every(Boolean)); + } finally { + globalThis.fetch = wrappedFetch; + invalidateQwenTokenPlanQuotaCache(connectionId); + } +}); + +test("agentrouter quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `agentrouter-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response(JSON.stringify({ data: { quota: 250_000 } }), { status: 200 }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchAgentrouterQuota(connectionId, { + providerSpecificData: { consoleApiKey: "system-access-token", newApiUserId: "42" }, + }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateAgentrouterQuotaCache(connectionId); + } +}); + +test("alibaba free tier quota reads inside the proxy context", async () => { + useAssignedProxy(); + const cookie = getAlibabaConsoleCookie({ alibabaConsoleCookie: "login_aliyunid_ticket=tick" }); + assert.ok(cookie); + const wrappedFetch = globalThis.fetch; + const contexts: boolean[] = []; + globalThis.fetch = (async () => { + contexts.push(hasAmbientProxyContext()); + return new Response( + JSON.stringify({ + code: "200", + data: { + DataV2: { + data: { + data: { + freeTierQuotas: [{ model: "qwen3-plus", freeTierOnly: true, quotaStatus: "VALID" }], + }, + }, + }, + }, + }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const entries = await fetchAlibabaFreeTierQuotaEntries( + { alibabaConsoleCookie: "login_aliyunid_ticket=tick" }, + `alibaba-proxy-${Date.now()}` + ); + assert.ok(entries); + assert.ok(contexts.length >= 1); + assert.ok(contexts.every(Boolean)); + } finally { + globalThis.fetch = wrappedFetch; + } +}); + +function bailianQuotaBody(used5h: number, usedWeekly: number): string { + return JSON.stringify({ + code: "Success", + data: { + codingPlanInstanceInfos: [ + { + planName: "Qwen3 Coder Next", + codingPlanQuotaInfo: { + per5HourUsedQuota: used5h, + per5HourTotalQuota: 100, + per5HourQuotaNextRefreshTime: 1718304000, + perWeekUsedQuota: usedWeekly, + perWeekTotalQuota: 100, + perWeekQuotaNextRefreshTime: 1718563200, + perBillMonthUsedQuota: 15, + perBillMonthTotalQuota: 100, + perBillMonthQuotaNextRefreshTime: 1719772800, + }, + }, + ], + }, + }); +} + +test("bailian quota retries the fallback host inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `bailian-proxy-${Date.now()}`; + const contexts: boolean[] = []; + const wrappedFetch = globalThis.fetch; + let calls = 0; + globalThis.fetch = (async () => { + calls += 1; + contexts.push(hasAmbientProxyContext()); + if (calls === 1) { + return new Response(JSON.stringify({ code: "ConsoleNeedLogin" }), { status: 200 }); + } + return new Response(bailianQuotaBody(30, 45), { status: 200 }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchBailianQuota(connectionId, { apiKey: "test-key" }); + assert.equal(calls, 2); + assert.ok(contexts.length === 2); + assert.ok(contexts.every(Boolean)); + assert.equal(quota?.percentUsed, 0.45); + } finally { + globalThis.fetch = wrappedFetch; + invalidateBailianQuotaCache(connectionId); + } +}); + +test("codex quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `codex-proxy-${Date.now()}`; + registerCodexConnection(connectionId, { accessToken: "access-token" }); + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ + rate_limit: { + primary_window: { used_percent: 80, reset_after_seconds: 60 }, + secondary_window: { used_percent: 40, reset_after_seconds: 120 }, + }, + }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const quota = await fetchCodexQuota(connectionId); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateCodexQuotaCache(connectionId); + unregisterCodexConnection(connectionId); + } +}); + +test("context7 quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `ctx7-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response(null, { + status: 200, + headers: { + "ratelimit-limit": "1000", + "ratelimit-remaining": "800", + "ratelimit-reset": "1790812800", + }, + }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchContext7Quota(connectionId, { apiKey: "ctx7sk-test-tok" }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateContext7QuotaCache(connectionId); + } +}); + +test("deepseek quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `deepseek-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ balance_infos: [{ currency: "USD", total_balance: "10.00" }] }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const quota = await fetchDeepseekQuota(connectionId, { apiKey: "test-key" }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateDeepseekQuotaCache(connectionId); + } +}); + +test("firecrawl quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `fc-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ + success: true, + data: { + remainingCredits: 700, + planCredits: 1000, + billingPeriodEnd: "2026-07-31T23:59:59Z", + }, + }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const quota = await fetchFirecrawlQuota(connectionId, { apiKey: "fc-test-key" }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateFirecrawlQuotaCache(connectionId); + } +}); + +function encodeVarint(value: number): Buffer { + const bytes: number[] = []; + let v = BigInt(value); + do { + let byte = Number(v & 0x7fn); + v >>= 7n; + if (v !== 0n) byte |= 0x80; + bytes.push(byte); + } while (v !== 0n); + return Buffer.from(bytes); +} + +function encodeTag(fieldNumber: number, wireType: number): Buffer { + return encodeVarint((fieldNumber << 3) | wireType); +} + +function encodeFixed32Field(fieldNumber: number, value: number): Buffer { + const body = Buffer.alloc(4); + body.writeFloatLE(value, 0); + return Buffer.concat([encodeTag(fieldNumber, 5), body]); +} + +function encodeLengthDelimited(fieldNumber: number, body: Buffer): Buffer { + return Buffer.concat([encodeTag(fieldNumber, 2), encodeVarint(body.length), body]); +} + +function encodeTimestampField(fieldNumber: number, seconds: number, nanos: number): Buffer { + const parts: Buffer[] = []; + if (seconds !== 0) { + parts.push(Buffer.concat([encodeTag(1, 0), encodeVarint(seconds)])); + } + if (nanos !== 0) { + parts.push(Buffer.concat([encodeTag(2, 0), encodeVarint(nanos)])); + } + return encodeLengthDelimited(fieldNumber, Buffer.concat(parts)); +} + +/** Framed gRPC-web GetGrokCreditsConfig response (same wire shape as the grok-cli test). */ +function grokCreditsFrame(usageRatio: number): ArrayBuffer { + const creditsInfo = Buffer.concat([ + encodeFixed32Field(1, usageRatio), + encodeTimestampField(5, 1784825940, 867850000), + ]); + const topMessage = encodeLengthDelimited(1, creditsInfo); + const header = Buffer.alloc(5); + header[0] = 0x00; + header.writeUInt32BE(topMessage.length, 1); + const framed = Buffer.concat([header, topMessage]); + return framed.buffer.slice(framed.byteOffset, framed.byteOffset + framed.byteLength); +} + +test("grok cli quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `grok-cli-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response(grokCreditsFrame(0.3), { status: 200 }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchGrokCliQuota(connectionId, { + credentials: { accessToken: "grok-token" }, + }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateGrokCliQuotaCache(connectionId); + } +}); + +function createGrokAuthFile(): string { + const dir = mkdtempSync(join(tmpdir(), "grok-auth-proxy-test-")); + const entry = { + key: "test-access-token", + refresh_token: "test-refresh-token", + expires_at: new Date(Date.now() + 3600_000).toISOString(), + email: "test@example.com", + auth_mode: "device_code", + oidc_issuer: "https://auth.x.ai", + oidc_client_id: "test-client-id", + }; + const authFile = join(dir, "auth.json"); + writeFileSync(authFile, JSON.stringify({ "test-entry": entry }, null, 2) + "\n"); + return authFile; +} + +test("grok web billing reads inside the proxy context", async () => { + useAssignedProxy(); + const authFile = createGrokAuthFile(); + process.env.GROK_AUTH_PATH = authFile; + const connectionId = `grok-proxy-${Date.now()}`; + const contexts: boolean[] = []; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async (input: RequestInfo | URL) => { + contexts.push(hasAmbientProxyContext()); + const urlStr = String(input); + if (urlStr.includes("billing")) { + return new Response( + JSON.stringify({ + config: { + creditUsagePercent: 15.3, + currentPeriod: { type: "WEEKLY", end: "2026-07-24T06:34:33.775Z" }, + productUsage: [{ product: "Api", usagePercent: 15 }], + onDemandCap: { val: 0 }, + onDemandUsed: { val: 0 }, + prepaidBalance: { val: 0 }, + }, + }), + { status: 200 } + ); + } + return new Response("not found", { status: 404 }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchGrokWebQuota(connectionId); + assert.ok(quota); + assert.ok(contexts.length >= 1); + assert.ok(contexts.every(Boolean)); + } finally { + globalThis.fetch = wrappedFetch; + invalidateGrokWebQuotaCache(connectionId); + try { + rmSync(authFile, { force: true }); + rmSync(join(authFile, ".."), { force: true }); + } catch { + // best effort cleanup + } + } +}); From 656ba2c5eeaa94b3a20cb19c1b1a3e072360ec46 Mon Sep 17 00:00:00 2001 From: Dizzle <112548150+maxmad64bis@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:47:33 +0200 Subject: [PATCH 2/3] fix(quota): route remaining console quota reads through the connection proxy Console quota reads used a direct fetch; they now run inside the assigned connection proxy context with no direct fallback. --- ...hers-remaining-through-connection-proxy.md | 1 + open-sse/services/jinaQuotaFetcher.ts | 4 +- open-sse/services/llmgatewayQuotaFetcher.ts | 6 +- open-sse/services/lyceumQuotaFetcher.ts | 7 +- open-sse/services/moonshotQuotaFetcher.ts | 4 +- .../services/newApiAggregatorQuotaFetcher.ts | 4 +- open-sse/services/opencodeQuotaFetcher.ts | 36 +- open-sse/services/openrouterQuotaFetcher.ts | 7 +- open-sse/services/tavilyQuotaFetcher.ts | 7 +- open-sse/services/v0QuotaFetcher.ts | 7 +- .../quota-fetchers-connection-proxy.test.ts | 343 ++++++++++++++++++ 11 files changed, 404 insertions(+), 22 deletions(-) create mode 100644 changelog.d/fixes/16344-quota-fetchers-remaining-through-connection-proxy.md diff --git a/changelog.d/fixes/16344-quota-fetchers-remaining-through-connection-proxy.md b/changelog.d/fixes/16344-quota-fetchers-remaining-through-connection-proxy.md new file mode 100644 index 000000000000..18d0f904846b --- /dev/null +++ b/changelog.d/fixes/16344-quota-fetchers-remaining-through-connection-proxy.md @@ -0,0 +1 @@ +- **fix(quota):** remaining console quota readings use the assigned connection egress path ([#16344](https://github.com/diegosouzapw/OmniRoute/pull/16344)) — thanks @maxmad64bis diff --git a/open-sse/services/jinaQuotaFetcher.ts b/open-sse/services/jinaQuotaFetcher.ts index 0e59daea4148..8aa617debcf8 100644 --- a/open-sse/services/jinaQuotaFetcher.ts +++ b/open-sse/services/jinaQuotaFetcher.ts @@ -1,6 +1,7 @@ import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; /** * Jina's wallet-balance endpoint only accepts the key as the `api_key` QUERY @@ -124,13 +125,14 @@ export async function fetchJinaQuota( await throttleQuotaFetch(); // Credential-bearing URL (see module header) — keep it local, never log it. const url = `${JINA_BALANCE_URL}?api_key=${encodeURIComponent(token)}`; - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers: { Accept: "application/json", }, signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), }); + if (!response) return null; if (!response.ok) { quotaCache.set(connectionId, { quota: null, fetchedAt: Date.now() }); diff --git a/open-sse/services/llmgatewayQuotaFetcher.ts b/open-sse/services/llmgatewayQuotaFetcher.ts index 5ede513fca45..5d589ff8f1de 100644 --- a/open-sse/services/llmgatewayQuotaFetcher.ts +++ b/open-sse/services/llmgatewayQuotaFetcher.ts @@ -48,6 +48,7 @@ import { } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; const LLMGATEWAY_CONFIG = { baseUrl: "https://api.llmgateway.io/v1", @@ -225,11 +226,12 @@ export function buildLlmgatewayQuota(key: LlmgatewayKeyFields): LlmgatewayQuota // ─── Core Fetcher ──────────────────────────────────────────────────────────── async function fetchJson( + connectionId: string, url: string, apiKey: string ): Promise<{ status: number; data: unknown } | null> { try { - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers: { Authorization: `Bearer ${apiKey}`, @@ -237,6 +239,7 @@ async function fetchJson( }, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; if (!response.ok) return { status: response.status, data: null }; const data = await response.json(); return { status: response.status, data }; @@ -272,6 +275,7 @@ export async function fetchLlmgatewayQuota( await throttleQuotaFetch(); const keyResult = await fetchJson( + connectionId, `${LLMGATEWAY_CONFIG.baseUrl}${LLMGATEWAY_CONFIG.keyPath}`, apiKey ); diff --git a/open-sse/services/lyceumQuotaFetcher.ts b/open-sse/services/lyceumQuotaFetcher.ts index ad23c19b2df0..29d2ec826de8 100644 --- a/open-sse/services/lyceumQuotaFetcher.ts +++ b/open-sse/services/lyceumQuotaFetcher.ts @@ -37,6 +37,7 @@ import { } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; const LYCEUM_CONFIG = { creditsUrl: "https://api.lyceum.technology/api/v2/external/billing/credits", @@ -153,11 +154,12 @@ export function buildLyceumQuota(fields: LyceumCreditsFields): LyceumQuota { // ─── Core Fetcher ──────────────────────────────────────────────────────────── async function fetchJson( + connectionId: string, url: string, apiKey: string ): Promise<{ status: number; data: unknown } | null> { try { - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers: { Authorization: `Bearer ${apiKey}`, @@ -165,6 +167,7 @@ async function fetchJson( }, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; if (!response.ok) return { status: response.status, data: null }; const data = await response.json(); return { status: response.status, data }; @@ -198,7 +201,7 @@ export async function fetchLyceumQuota( try { await throttleQuotaFetch(); - const result = await fetchJson(LYCEUM_CONFIG.creditsUrl, apiKey); + const result = await fetchJson(connectionId, LYCEUM_CONFIG.creditsUrl, apiKey); if (!result || result.status !== 200) { quotaCache.delete(connectionId); return null; diff --git a/open-sse/services/moonshotQuotaFetcher.ts b/open-sse/services/moonshotQuotaFetcher.ts index 5f8acee2c879..b63359fc8fe6 100644 --- a/open-sse/services/moonshotQuotaFetcher.ts +++ b/open-sse/services/moonshotQuotaFetcher.ts @@ -14,6 +14,7 @@ import { toNumber } from "@/shared/utils/numeric"; import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { isUrlHostOrSubdomain } from "../utils/urlHost.ts"; import { isKimiCodingBaseUrl, @@ -171,7 +172,7 @@ export async function fetchMoonshotQuota( try { await throttleQuotaFetch(); - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers: { Authorization: authHeader, @@ -180,6 +181,7 @@ export async function fetchMoonshotQuota( }, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; if (response.status === 401 || response.status === 403) { quotaCache.delete(connectionId); diff --git a/open-sse/services/newApiAggregatorQuotaFetcher.ts b/open-sse/services/newApiAggregatorQuotaFetcher.ts index cdcfca32744e..9c261d5fc090 100644 --- a/open-sse/services/newApiAggregatorQuotaFetcher.ts +++ b/open-sse/services/newApiAggregatorQuotaFetcher.ts @@ -33,6 +33,7 @@ import type { QuotaInfo } from "./quotaPreflight.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { toNumber } from "@/shared/utils/numeric"; const SELF_PATH = "/api/user/self"; @@ -174,7 +175,7 @@ export async function fetchNewApiAggregatorQuota( try { await throttleQuotaFetch(); - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers: { Authorization: `Bearer ${systemAccessToken}`, @@ -184,6 +185,7 @@ export async function fetchNewApiAggregatorQuota( }, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; if (response.status === 401 || response.status === 403) { quotaCache.delete(connectionId); diff --git a/open-sse/services/opencodeQuotaFetcher.ts b/open-sse/services/opencodeQuotaFetcher.ts index f49a14cd8ccd..b1fbe3ed7c51 100644 --- a/open-sse/services/opencodeQuotaFetcher.ts +++ b/open-sse/services/opencodeQuotaFetcher.ts @@ -12,6 +12,7 @@ import { registerQuotaFetcher, registerQuotaWindows, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { buildOpencodeBackgroundHeaders } from "../utils/opencodeHeaders.ts"; const OPENCODE_QUOTA_URL = @@ -41,6 +42,15 @@ interface CacheEntry { // In-memory cache: connectionId → successful quota bound to its normalized API key const quotaCache = new Map(); +function normalizeApiKey(connection?: Record): string { + return typeof connection?.apiKey === "string" + ? connection.apiKey + .trim() + .replace(/^Bearer\s+/i, "") + .trim() + : ""; +} + // One-time 404 warning per URL (avoids spamming on every request) const _warned404Urls = new Set(); @@ -160,23 +170,21 @@ export async function fetchOpencodeQuota( connectionId: string, connection?: Record ): Promise { - const apiKey = - typeof connection?.apiKey === "string" - ? connection.apiKey - .trim() - .replace(/^Bearer\s+/i, "") - .trim() - : ""; + // The dashboard usage path calls without a stored connection identifier; keep + // the previous behaviour there by resolving the shared no-auth fallback + // instead of reporting an unknown reading. + const effectiveConnectionId = connectionId || "noauth"; + const apiKey = normalizeApiKey(connection); if (!apiKey) return null; - const cached = quotaCache.get(connectionId); + const cached = quotaCache.get(effectiveConnectionId); if (cached && cached.apiKey === apiKey && Date.now() - cached.fetchedAt < CACHE_TTL_MS) { return cached.quota; } try { await throttleQuotaFetch(); - const response = await fetch(OPENCODE_QUOTA_URL, { + const response = await fetchWithConnectionProxy(effectiveConnectionId, OPENCODE_QUOTA_URL, { method: "GET", headers: { Authorization: `Bearer ${apiKey}`, @@ -186,10 +194,11 @@ export async function fetchOpencodeQuota( // ("Bun fetch") without x-opencode-session are flagged, with hard // errors announced from 2026-09-06. Send the CLI identity + a stable // per-connection session fingerprint (see buildOpencodeBackgroundHeaders). - ...buildOpencodeBackgroundHeaders({ seed: connectionId }), + ...buildOpencodeBackgroundHeaders({ seed: effectiveConnectionId }), }, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; if (!response.ok) { if (response.status === 404 && !_warned404Urls.has(OPENCODE_QUOTA_URL)) { @@ -199,7 +208,8 @@ export async function fetchOpencodeQuota( "Verify OMNIROUTE_OPENCODE_QUOTA_URL when using a relay or test server." ); } - if (response.status === 401 || response.status === 403) quotaCache.delete(connectionId); + if (response.status === 401 || response.status === 403) + quotaCache.delete(effectiveConnectionId); return null; } @@ -213,7 +223,7 @@ export async function fetchOpencodeQuota( const quota = parseOpencodeQuotaResponse(data); if (!quota) return null; - quotaCache.set(connectionId, { quota, fetchedAt: Date.now(), apiKey }); + quotaCache.set(effectiveConnectionId, { quota, fetchedAt: Date.now(), apiKey }); return quota; } catch { return null; @@ -226,7 +236,7 @@ export async function fetchOpencodeQuota( * Force-invalidate the cache for a connection (e.g., after receiving quota headers). */ export function invalidateOpencodeQuotaCache(connectionId: string): void { - quotaCache.delete(connectionId); + quotaCache.delete(connectionId || "noauth"); } // ─── Registration ───────────────────────────────────────────────────────────── diff --git a/open-sse/services/openrouterQuotaFetcher.ts b/open-sse/services/openrouterQuotaFetcher.ts index 2fc7a94422c0..5ebd442602a1 100644 --- a/open-sse/services/openrouterQuotaFetcher.ts +++ b/open-sse/services/openrouterQuotaFetcher.ts @@ -36,6 +36,7 @@ import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { getFreeWindowStatus, isFreeVariantModel, @@ -283,11 +284,12 @@ function checkFreeWindowExhausted( // ─── Core Fetcher ──────────────────────────────────────────────────────────── async function fetchJson( + connectionId: string, url: string, apiKey: string ): Promise<{ status: number; data: unknown } | null> { try { - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers: { Authorization: `Bearer ${apiKey}`, @@ -295,6 +297,7 @@ async function fetchJson( }, signal: AbortSignal.timeout(8_000), }); + if (!response) return null; if (!response.ok) return { status: response.status, data: null }; const data = await response.json(); return { status: response.status, data }; @@ -385,10 +388,12 @@ export async function fetchOpenrouterQuota( await throttleQuotaFetch(); const keyResult = await fetchJson( + connectionId, `${OPENROUTER_CONFIG.baseUrl}${OPENROUTER_CONFIG.keyPath}`, apiKey ); const creditsResult = await fetchJson( + connectionId, `${OPENROUTER_CONFIG.baseUrl}${OPENROUTER_CONFIG.creditsPath}`, apiKey ); diff --git a/open-sse/services/tavilyQuotaFetcher.ts b/open-sse/services/tavilyQuotaFetcher.ts index 3b94d7d3beaa..511e392f9418 100644 --- a/open-sse/services/tavilyQuotaFetcher.ts +++ b/open-sse/services/tavilyQuotaFetcher.ts @@ -20,6 +20,7 @@ import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; import { isUrlHost } from "../utils/urlHost.ts"; import { toNumberOrNull } from "@/shared/utils/numeric"; @@ -180,7 +181,7 @@ export async function fetchTavilyQuota( try { await throttleQuotaFetch(); - const response = await fetch(TAVILY_USAGE_URL, { + const response = await fetchWithConnectionProxy(connectionId, TAVILY_USAGE_URL, { method: "GET", headers: { Authorization: `Bearer ${token}`, @@ -188,6 +189,10 @@ export async function fetchTavilyQuota( }, signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), }); + if (!response) { + quotaCache.set(connectionId, { quota: null, fetchedAt: Date.now() }); + return null; + } if (!response.ok) { quotaCache.set(connectionId, { quota: null, fetchedAt: Date.now() }); diff --git a/open-sse/services/v0QuotaFetcher.ts b/open-sse/services/v0QuotaFetcher.ts index 853b67580e3d..3bc0b3b5ab85 100644 --- a/open-sse/services/v0QuotaFetcher.ts +++ b/open-sse/services/v0QuotaFetcher.ts @@ -22,6 +22,7 @@ import { registerQuotaFetcher, registerQuotaWindows, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; import { throttleQuotaFetch } from "./quotaFetchThrottle.ts"; +import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts"; const V0_CONFIG = { baseUrl: "https://api.v0.dev", @@ -115,17 +116,19 @@ interface WindowFetchResult { * rate-limits (dailyOps) calls — the only difference is the response shape parser. */ async function fetchWindow( + connectionId: string, url: string, headers: Record, parse: (data: unknown) => { billingType: string | null; window: WindowFetchResult["window"] } ): Promise { try { await throttleQuotaFetch(); - const response = await fetch(url, { + const response = await fetchWithConnectionProxy(connectionId, url, { method: "GET", headers, signal: AbortSignal.timeout(8_000), }); + if (!response) return { window: null, billingType: null, invalidCredential: false }; if (response.status === 401 || response.status === 403) { return { window: null, billingType: null, invalidCredential: true }; @@ -219,6 +222,7 @@ export async function fetchV0Quota( let billingType = "unknown"; const billingResult = await fetchWindow( + connectionId, `${V0_CONFIG.baseUrl}${V0_CONFIG.billingPath}`, headers, parseBillingResponse @@ -229,6 +233,7 @@ export async function fetchV0Quota( } const rateLimitsResult = await fetchWindow( + connectionId, `${V0_CONFIG.baseUrl}${V0_CONFIG.rateLimitsPath}`, headers, parseRateLimitsResponse diff --git a/tests/unit/quota-fetchers-connection-proxy.test.ts b/tests/unit/quota-fetchers-connection-proxy.test.ts index 85c30d2b3de5..7231c503d272 100644 --- a/tests/unit/quota-fetchers-connection-proxy.test.ts +++ b/tests/unit/quota-fetchers-connection-proxy.test.ts @@ -60,6 +60,40 @@ import { fetchQwenTokenPlanQuota, invalidateQwenTokenPlanQuotaCache, } from "../../open-sse/services/qwenTokenPlanQuotaFetcher.ts"; +import { + fetchJinaQuota, + invalidateJinaQuotaCache, +} from "../../open-sse/services/jinaQuotaFetcher.ts"; +import { + fetchLlmgatewayQuota, + invalidateLlmgatewayQuotaCache, +} from "../../open-sse/services/llmgatewayQuotaFetcher.ts"; +import { + fetchLyceumQuota, + invalidateLyceumQuotaCache, +} from "../../open-sse/services/lyceumQuotaFetcher.ts"; +import { + fetchMoonshotQuota, + invalidateMoonshotQuotaCache, +} from "../../open-sse/services/moonshotQuotaFetcher.ts"; +import { + fetchNewApiAggregatorQuota, + invalidateNewApiAggregatorQuotaCache, +} from "../../open-sse/services/newApiAggregatorQuotaFetcher.ts"; +import { + fetchOpencodeQuota, + invalidateOpencodeQuotaCache, +} from "../../open-sse/services/opencodeQuotaFetcher.ts"; +import { + fetchOpenrouterQuota, + invalidateOpenrouterQuotaCache, +} from "../../open-sse/services/openrouterQuotaFetcher.ts"; +import { + fetchTavilyQuota, + invalidateTavilyQuotaCache, +} from "../../open-sse/services/tavilyQuotaFetcher.ts"; +import { fetchV0Quota, invalidateV0QuotaCache } from "../../open-sse/services/v0QuotaFetcher.ts"; +import { readFileSync } from "node:fs"; // socks5 skips the TCP reachability probe in runWithProxyContext, so the mock // fetch below never races the fast-fail path; the proxy context is still set. @@ -578,3 +612,312 @@ test("grok web billing reads inside the proxy context", async () => { } } }); + +test("jina quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `jina-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ + code: 200, + data: { user_id: "u1", wallet: { total_balance: 5000000 } }, + }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const quota = await fetchJinaQuota(connectionId, { apiKey: "jina-tok" }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateJinaQuotaCache(connectionId); + } +}); + +test("llmgateway quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `llmgw-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ + data: { + devPlan: "pro", + devPlanCreditsUsed: "25", + devPlanCreditsLimit: "237", + devPlanCreditsRemaining: "212.00", + devPlanPremiumWeeklyLimit: "35.55", + devPlanPremiumCreditsUsed: "5.00", + devPlanPremiumWeekResetsAt: "2026-08-28T12:00:00.000Z", + }, + }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const quota = await fetchLlmgatewayQuota(connectionId, { apiKey: "llmgtwy_test" }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateLlmgatewayQuotaCache(connectionId); + } +}); + +test("lyceum quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `lyceum-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ available_credits: 50, used_credits: 10, remaining_credits: 40 }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const quota = await fetchLyceumQuota(connectionId, { apiKey: "lk_test" }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateLyceumQuotaCache(connectionId); + } +}); + +test("moonshot quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `ms-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ + code: 0, + data: { available_balance: 2.5, voucher_balance: 0, cash_balance: 2.5 }, + status: true, + }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const quota = await fetchMoonshotQuota(connectionId, { + apiKey: "sk-test", + providerSpecificData: { baseUrl: "https://api.moonshot.cn/v1" }, + }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateMoonshotQuotaCache(connectionId); + } +}); + +test("new-api aggregator quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `agg-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response(JSON.stringify({ data: { quota: 250_000 } }), { status: 200 }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchNewApiAggregatorQuota(connectionId, { + providerSpecificData: { + consoleApiKey: "system-access-token", + newApiUserId: "42", + newApiAggregatorBalance: true, + baseUrl: "https://my-newapi.example.com/v1", + }, + }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateNewApiAggregatorQuotaCache(connectionId); + } +}); + +test("opencode quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `opencode-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ + usage: { + rolling: { status: "ok", percent: 25, resetsAt: "2026-09-01T01:02:03.000Z" }, + weekly: { status: "ok", percent: 50, resetsAt: "2026-09-05T04:05:06.000Z" }, + monthly: { status: "ok", percent: 10, resetsAt: "2026-09-30T07:08:09.000Z" }, + }, + }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const quota = await fetchOpencodeQuota(connectionId, { apiKey: "opencode-key" }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateOpencodeQuotaCache(connectionId); + } +}); + +test("opencode quota invalidation clears the shared no-auth reading", async () => { + useAssignedProxy(); + let calls = 0; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + calls += 1; + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ + usage: { + rolling: { status: "ok", percent: 25, resetsAt: "2026-09-01T01:02:03.000Z" }, + weekly: { status: "ok", percent: 50, resetsAt: "2026-09-05T04:05:06.000Z" }, + monthly: { status: "ok", percent: 10, resetsAt: "2026-09-30T07:08:09.000Z" }, + }, + }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const first = await fetchOpencodeQuota("", { apiKey: "opencode-key" }); + assert.ok(first); + assert.equal(observed, true); + assert.equal(calls, 1); + invalidateOpencodeQuotaCache(""); + const second = await fetchOpencodeQuota("", { apiKey: "opencode-key" }); + assert.ok(second); + assert.equal(calls, 2); + } finally { + globalThis.fetch = wrappedFetch; + invalidateOpencodeQuotaCache(""); + invalidateOpencodeQuotaCache("noauth"); + } +}); + +test("openrouter quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `openrouter-proxy-${Date.now()}`; + const contexts: boolean[] = []; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async (url: RequestInfo | URL) => { + contexts.push(hasAmbientProxyContext()); + if (String(url).endsWith("/key")) { + return new Response( + JSON.stringify({ + data: { limit: 100, limit_remaining: 10, limit_reset: null, is_free_tier: false }, + }), + { status: 200 } + ); + } + return new Response(JSON.stringify({ data: { total_credits: 100, total_usage: 90 } }), { + status: 200, + }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchOpenrouterQuota(connectionId, { apiKey: "test-key" }); + assert.ok(quota); + assert.ok(contexts.length === 2); + assert.ok(contexts.every(Boolean)); + } finally { + globalThis.fetch = wrappedFetch; + invalidateOpenrouterQuotaCache(connectionId); + } +}); + +test("tavily quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `tavily-proxy-${Date.now()}`; + let observed: boolean | null = null; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + observed = hasAmbientProxyContext(); + return new Response( + JSON.stringify({ + account: { current_plan: "Researcher", plan_usage: 10, plan_limit: 1000 }, + }), + { status: 200 } + ); + }) as typeof globalThis.fetch; + try { + const quota = await fetchTavilyQuota(connectionId, { apiKey: "mock-token" }); + assert.ok(quota); + assert.equal(observed, true); + } finally { + globalThis.fetch = wrappedFetch; + invalidateTavilyQuotaCache(connectionId); + } +}); + +test("v0 quota reads inside the proxy context", async () => { + useAssignedProxy(); + const connectionId = `v0-proxy-${Date.now()}`; + const contexts: boolean[] = []; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async (url: RequestInfo | URL) => { + contexts.push(hasAmbientProxyContext()); + if (String(url).endsWith("/v1/user/billing")) { + return new Response( + JSON.stringify({ + billingType: "credit", + data: { remaining: 20, limit: 100, reset: "2026-08-01T00:00:00.000Z" }, + }), + { status: 200 } + ); + } + return new Response(JSON.stringify({ remaining: 900, limit: 1000 }), { status: 200 }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchV0Quota(connectionId, { apiKey: "v0-key" }); + assert.ok(quota); + assert.ok(contexts.length === 2); + assert.ok(contexts.every(Boolean)); + } finally { + globalThis.fetch = wrappedFetch; + invalidateV0QuotaCache(connectionId); + } +}); + +test("every remaining quota fetcher routes through the connection proxy helper", async () => { + const { readdirSync } = await import("node:fs"); + const { dirname } = await import("node:path"); + const { fileURLToPath } = await import("node:url"); + const here = dirname(fileURLToPath(import.meta.url)); + const dir = join(here, "..", "..", "open-sse", "services"); + const targets = [ + "jinaQuotaFetcher.ts", + "llmgatewayQuotaFetcher.ts", + "lyceumQuotaFetcher.ts", + "moonshotQuotaFetcher.ts", + "newApiAggregatorQuotaFetcher.ts", + "opencodeQuotaFetcher.ts", + "openrouterQuotaFetcher.ts", + "tavilyQuotaFetcher.ts", + "v0QuotaFetcher.ts", + ]; + const onDisk = new Set(readdirSync(dir)); + const offenders: string[] = []; + for (const name of targets) { + assert.ok(onDisk.has(name), `missing fetcher file ${name}`); + const source = readFileSync(join(dir, name), "utf8"); + const direct = [...source.matchAll(/await fetch\(/g)].length; + const helper = [...source.matchAll(/fetchWithConnectionProxy\(/g)].length; + if (direct > 0 || helper < 1) offenders.push(`${name} (direct=${direct} helper=${helper})`); + } + assert.deepEqual(offenders, []); +}); From 858ef2f045d2be6555ccc6e293ec23e50ccb4547 Mon Sep 17 00:00:00 2001 From: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:44:25 -0300 Subject: [PATCH 3/3] fix(quota): drop the stale opencode no-auth key and make grok pass its connection id - opencodeQuotaFetcher: keyless OpenCode Free is gone (#16310), so the 'noauth' fallback id resolved a proxy for a connection that cannot exist and shared one cache entry across every empty-id reading. Key the cache, the session seed and the proxy lookup on the real connection id. - grokQuotaFetcher: the billing/OIDC helpers took an optional id and passed 'connectionId ?? ""' to fetchWithConnectionProxy, which reports an unknown reading for an empty id. Every caller already has the id (fetchGrokWebQuota, xaiOauthQuotaFetcher), so make it required. --- open-sse/services/grokQuotaFetcher.ts | 21 ++++------ open-sse/services/opencodeQuotaFetcher.ts | 17 +++----- .../quota-fetchers-connection-proxy.test.ts | 42 ++++++++++++++----- 3 files changed, 46 insertions(+), 34 deletions(-) diff --git a/open-sse/services/grokQuotaFetcher.ts b/open-sse/services/grokQuotaFetcher.ts index 6cfaddf44165..004b7d8dfa2e 100644 --- a/open-sse/services/grokQuotaFetcher.ts +++ b/open-sse/services/grokQuotaFetcher.ts @@ -208,13 +208,13 @@ function writeRefreshedTokens( async function discoverTokenEndpoint( issuer: string, signal: AbortSignal, - connectionId?: string | null + connectionId: string ): Promise { const discoveryUrl = `${issuer.replace(/\/$/, "")}/.well-known/openid-configuration`; if (!isAllowedXaiUrl(discoveryUrl)) { throw new Error("invalid oidc issuer"); } - const res = await fetchWithConnectionProxy(connectionId ?? "", discoveryUrl, { + const res = await fetchWithConnectionProxy(connectionId, discoveryUrl, { headers: { Accept: "application/json" }, signal, }); @@ -229,7 +229,7 @@ async function discoverTokenEndpoint( async function refreshAccessToken( auth: ResolvedAuth, signal: AbortSignal, - connectionId?: string | null + connectionId: string ): Promise { if (!auth.refreshToken) { throw new Error("auth expired — run `grok login`"); @@ -239,7 +239,7 @@ async function refreshAccessToken( } const tokenEndpoint = await discoverTokenEndpoint(auth.issuer, signal, connectionId); - const res = await fetchWithConnectionProxy(connectionId ?? "", tokenEndpoint, { + const res = await fetchWithConnectionProxy(connectionId, tokenEndpoint, { method: "POST", headers: { "Content-Type": "application/x-www-form-urlencoded", @@ -297,10 +297,7 @@ function needsRefresh(auth: ResolvedAuth): boolean { return auth.expiresAtMs <= Date.now() + EXPIRY_SKEW_MS; } -async function resolveAuth( - signal: AbortSignal, - connectionId?: string | null -): Promise { +async function resolveAuth(signal: AbortSignal, connectionId: string): Promise { const file = readAuthFile(); if (!file) throw new Error("no grok auth — run `grok login`"); const auth = pickAuthEntry(file); @@ -343,8 +340,8 @@ function resetLocalLabel(endIso?: string): string { */ export async function fetchGrokBillingWithToken( token: string, - signal?: AbortSignal, - connectionId?: string | null + signal: AbortSignal | undefined, + connectionId: string ): Promise { const controller = signal ? null : new AbortController(); const timeout = controller ? setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS) : null; @@ -366,9 +363,9 @@ export function grokBillingSnapshotToQuotaInfo(snap: GrokBillingSnapshot): Quota async function fetchBilling( token: string, signal: AbortSignal, - connectionId?: string | null + connectionId: string ): Promise { - const res = await fetchWithConnectionProxy(connectionId ?? "", BILLING_URL, { + const res = await fetchWithConnectionProxy(connectionId, BILLING_URL, { method: "GET", headers: { Authorization: `Bearer ${token}`, diff --git a/open-sse/services/opencodeQuotaFetcher.ts b/open-sse/services/opencodeQuotaFetcher.ts index b8e47454b16c..9e56c7678674 100644 --- a/open-sse/services/opencodeQuotaFetcher.ts +++ b/open-sse/services/opencodeQuotaFetcher.ts @@ -170,21 +170,17 @@ export async function fetchOpencodeQuota( connectionId: string, connection?: Record ): Promise { - // The dashboard usage path calls without a stored connection identifier; keep - // the previous behaviour there by resolving the shared no-auth fallback - // instead of reporting an unknown reading. - const effectiveConnectionId = connectionId || "noauth"; const apiKey = normalizeApiKey(connection); if (!apiKey) return null; - const cached = quotaCache.get(effectiveConnectionId); + const cached = quotaCache.get(connectionId); if (cached && cached.apiKey === apiKey && Date.now() - cached.fetchedAt < CACHE_TTL_MS) { return cached.quota; } try { await throttleQuotaFetch(); - const response = await fetchWithConnectionProxy(effectiveConnectionId, OPENCODE_QUOTA_URL, { + const response = await fetchWithConnectionProxy(connectionId, OPENCODE_QUOTA_URL, { method: "GET", headers: { Authorization: `Bearer ${apiKey}`, @@ -194,7 +190,7 @@ export async function fetchOpencodeQuota( // ("Bun fetch") without x-opencode-session are flagged, with hard // errors announced from 2026-09-06. Send the CLI identity + a stable // per-connection session fingerprint (see buildOpencodeBackgroundHeaders). - ...buildOpencodeBackgroundHeaders({ seed: effectiveConnectionId }), + ...buildOpencodeBackgroundHeaders({ seed: connectionId }), }, signal: AbortSignal.timeout(8_000), }); @@ -208,8 +204,7 @@ export async function fetchOpencodeQuota( "Verify OMNIROUTE_OPENCODE_QUOTA_URL when using a relay or test server." ); } - if (response.status === 401 || response.status === 403) - quotaCache.delete(effectiveConnectionId); + if (response.status === 401 || response.status === 403) quotaCache.delete(connectionId); return null; } @@ -223,7 +218,7 @@ export async function fetchOpencodeQuota( const quota = parseOpencodeQuotaResponse(data); if (!quota) return null; - quotaCache.set(effectiveConnectionId, { quota, fetchedAt: Date.now(), apiKey }); + quotaCache.set(connectionId, { quota, fetchedAt: Date.now(), apiKey }); return quota; } catch { return null; @@ -236,7 +231,7 @@ export async function fetchOpencodeQuota( * Force-invalidate the cache for a connection (e.g., after receiving quota headers). */ export function invalidateOpencodeQuotaCache(connectionId: string): void { - quotaCache.delete(connectionId || "noauth"); + quotaCache.delete(connectionId); } // ─── Registration ───────────────────────────────────────────────────────────── diff --git a/tests/unit/quota-fetchers-connection-proxy.test.ts b/tests/unit/quota-fetchers-connection-proxy.test.ts index 7231c503d272..4abc109f6548 100644 --- a/tests/unit/quota-fetchers-connection-proxy.test.ts +++ b/tests/unit/quota-fetchers-connection-proxy.test.ts @@ -775,14 +775,34 @@ test("opencode quota reads inside the proxy context", async () => { } }); -test("opencode quota invalidation clears the shared no-auth reading", async () => { +test("opencode quota without a connection identifier makes no fetch outside the proxy", async () => { + // Keyless OpenCode Free is gone (#16310): every OpenCode reading belongs to a + // stored connection, so there is no shared "noauth" reading to resolve a proxy + // for. An empty id reports no reading, like every other quota fetcher. useAssignedProxy(); let calls = 0; - let observed: boolean | null = null; const wrappedFetch = globalThis.fetch; globalThis.fetch = (async () => { calls += 1; - observed = hasAmbientProxyContext(); + return new Response("{}", { status: 200 }); + }) as typeof globalThis.fetch; + try { + const quota = await fetchOpencodeQuota("", { apiKey: "opencode-key" }); + assert.equal(quota, null); + assert.equal(calls, 0); + } finally { + globalThis.fetch = wrappedFetch; + invalidateOpencodeQuotaCache(""); + } +}); + +test("opencode quota caches and invalidates by the connection id, not a shared no-auth key", async () => { + useAssignedProxy(); + const connectionId = `opencode-cache-${Date.now()}`; + let calls = 0; + const wrappedFetch = globalThis.fetch; + globalThis.fetch = (async () => { + calls += 1; return new Response( JSON.stringify({ usage: { @@ -795,18 +815,18 @@ test("opencode quota invalidation clears the shared no-auth reading", async () = ); }) as typeof globalThis.fetch; try { - const first = await fetchOpencodeQuota("", { apiKey: "opencode-key" }); - assert.ok(first); - assert.equal(observed, true); + assert.ok(await fetchOpencodeQuota(connectionId, { apiKey: "opencode-key" })); + assert.ok(await fetchOpencodeQuota(connectionId, { apiKey: "opencode-key" })); assert.equal(calls, 1); - invalidateOpencodeQuotaCache(""); - const second = await fetchOpencodeQuota("", { apiKey: "opencode-key" }); - assert.ok(second); + invalidateOpencodeQuotaCache("noauth"); + assert.ok(await fetchOpencodeQuota(connectionId, { apiKey: "opencode-key" })); + assert.equal(calls, 1); + invalidateOpencodeQuotaCache(connectionId); + assert.ok(await fetchOpencodeQuota(connectionId, { apiKey: "opencode-key" })); assert.equal(calls, 2); } finally { globalThis.fetch = wrappedFetch; - invalidateOpencodeQuotaCache(""); - invalidateOpencodeQuotaCache("noauth"); + invalidateOpencodeQuotaCache(connectionId); } });