fix(usage): route usage readings through the assigned connection proxy

changelog.d/fixes/16773-usage-readers-connection-proxy.md

@@ -0,0 +1 @@
+- **fix(usage):** route usage readings through the assigned connection proxy ([#16773](https://github.com/diegosouzapw/OmniRoute/pull/16773)) — thanks @maxmad64bis

open-sse/services/connectionProxyFetch.ts

@@ -0,0 +1,113 @@
+/**
+ * connectionProxyFetch.ts — run quota fetcher requests through the proxy
+ * assigned to the connection.
+ *
+ * Quota fetchers call provider consoles directly with the global fetch. On an
+ * installation behind a proxy, that reads the quota of the wrong egress
+ * address and exposes the origin. This helper resolves the proxy bound to the
+ * connection and runs the fetch inside that proxy context, so the reading
+ * uses the assigned egress path. Without an assigned proxy the fetch runs
+ * unchanged. A proxy or resolution failure returns null so the reading is
+ * reported, never retried outside the proxy.
+ */
+
+// NOTE: proxyFetch and error are imported lazily inside the functions below
+// (not statically at the top level) so that bundling this module (e.g. the
+// MCP server bundle) does not initialize the proxy transport at load time.
+
+type ProxyValue = unknown;
+
+export interface ProxyResolver {
+  resolveProxyForConnection: (connectionId: string) => Promise<{ proxy?: ProxyValue } | null>;
+}
+
+interface ConnectionProxyFetchDeps {
+  resolveProxyForConnection?: ProxyResolver["resolveProxyForConnection"];
+  fetchFn?: typeof globalThis.fetch;
+}
+
+async function defaultResolveProxyForConnection(
+  connectionId: string
+): Promise<{ proxy?: ProxyValue } | null> {
+  const override = connectionProxyTestOverride.resolveProxyForConnection;
+  if (override) return override(connectionId);
+  const { resolveProxyForConnection } = await import("@/lib/db/settings");
+  return resolveProxyForConnection(connectionId);
+}
+
+/**
+ * Test-only override for the connection proxy resolution (avoids seeding the
+ * database in unit tests). Production code never sets it.
+ */
+export const connectionProxyTestOverride: {
+  resolveProxyForConnection:
+    ((connectionId: string) => Promise<{ proxy?: ProxyValue } | null>) | null;
+} = { resolveProxyForConnection: null };
+
+/**
+ * Resolve the proxy assigned to a quota connection. Null means no proxy is
+ * assigned and the fetch runs unchanged.
+ */
+export async function resolveConnectionProxy(
+  connectionId: string,
+  deps: ConnectionProxyFetchDeps = {}
+): Promise<ProxyValue | null> {
+  const resolve = deps.resolveProxyForConnection ?? defaultResolveProxyForConnection;
+  const resolved = await resolve(connectionId);
+  return resolved?.proxy ?? null;
+}
+
+/**
+ * Fetch a quota URL through the connection proxy context. Returns the
+ * response, or null when no proxy is assigned, when the connection has no
+ * identifier to resolve a proxy for, and when the proxied fetch fails — the
+ * caller reports the missing reading instead of refetching outside the proxy.
+ */
+export async function fetchWithConnectionProxy(
+  connectionId: string,
+  url: string,
+  init?: RequestInit,
+  deps: ConnectionProxyFetchDeps = {}
+): Promise<Response | null> {
+  const fetchFn = deps.fetchFn ?? globalThis.fetch;
+  const { runWithProxyContext } = await import("../utils/proxyFetch.ts");
+  const { sanitizeErrorMessage } = await import("../utils/error.ts");
+  // Importing the proxy transport replaces globalThis.fetch with its own
+  // interceptor as a side effect. When the caller provided the fetch (a test
+  // mock, or any fetch installed before this call), put it back so the next
+  // request of the same quota reading keeps using it instead of the
+  // interceptor. In production the global fetch is already the interceptor
+  // (installed once at startup), so this is a no-op there.
+  if (deps.fetchFn === undefined && globalThis.fetch !== fetchFn) {
+    globalThis.fetch = fetchFn;
+  }
+  if (!connectionId) {
+    console.warn(
+      "[ConnectionProxyFetch] Missing connection identifier: reporting an unknown reading instead of fetching outside a proxy."
+    );
+    return null;
+  }
+  let proxy: ProxyValue | null;
+  try {
+    proxy = await resolveConnectionProxy(connectionId, deps);
+  } catch (error) {
+    console.warn(
+      `[ConnectionProxyFetch] Proxy resolution failed for ${connectionId}: ${sanitizeErrorMessage(
+        error instanceof Error ? error.message : String(error)
+      )}`
+    );
+    return null;
+  }
+  const task = () => fetchFn(url, init);
+  if (!proxy) return task();
+  try {
+    return await runWithProxyContext(proxy, task);
+  } catch (error) {
+    console.warn(
+      `[ConnectionProxyFetch] Proxied quota fetch failed for ${connectionId}: ${sanitizeErrorMessage(
+        error instanceof Error ? error.message : String(error)
+      )}`
+    );
+    return null;
+  }
+}

open-sse/services/conolUsage.ts

@@ -1,3 +1,4 @@
+import { fetchWithConnectionProxy } from "./connectionProxyFetch.ts";
 import { normalizeConolCookie } from "./conolAuth.ts";
 
 interface UsageQuota {
@@ -77,7 +78,8 @@ function readProviderValue(data: unknown, keys: readonly string[]): string {
 
 export async function getConolUsage(
   apiKey: unknown,
-  providerSpecificData?: unknown
+  providerSpecificData?: unknown,
+  connectionId?: string
 ): Promise<ConolUsageResult | { message: string }> {
   const raw =
     readProviderValue(providerSpecificData, [
@@ -89,15 +91,30 @@ export async function getConolUsage(
   if (!cookie) return { message: "Missing Conol session cookie" };
 
   try {
-    const response = await fetch("https://conol.ai/api/billing/balance", {
-      method: "GET",
-      headers: {
-        accept: "application/json",
-        cookie,
-        referer: "https://conol.ai/home",
-      },
-      signal: AbortSignal.timeout(15_000),
-    });
+    const response = connectionId
+      ? await fetchWithConnectionProxy(
+          connectionId,
+          "https://conol.ai/api/billing/balance",
+          {
+            method: "GET",
+            headers: {
+              accept: "application/json",
+              cookie,
+              referer: "https://conol.ai/home",
+            },
+            signal: AbortSignal.timeout(15_000),
+          }
+        )
+      : await globalThis.fetch("https://conol.ai/api/billing/balance", {
+          method: "GET",
+          headers: {
+            accept: "application/json",
+            cookie,
+            referer: "https://conol.ai/home",
+          },
+          signal: AbortSignal.timeout(15_000),
+        });
+    if (!response) return { message: "Conol balance request failed" };
     if (response.status === 401 || response.status === 403) {
       return { message: "Conol session expired or is invalid" };
     }

open-sse/services/usage.ts

@@ -148,7 +148,7 @@ export async function getUsageForProvider(
 
   switch (provider) {
     case "github":
-      return await getGitHubUsage(accessToken, providerSpecificData);
+      return await getGitHubUsage(accessToken, providerSpecificData, id || "");
     case "antigravity":
     case "agy":
       return await getAntigravityUsage(
@@ -160,16 +160,18 @@ export async function getUsageForProvider(
         options
       );
     case "claude":
-      return await getClaudeUsage(accessToken, { forceRefresh: options.forceRefresh === true });
+      return await getClaudeUsage(accessToken, id || "", {
+        forceRefresh: options.forceRefresh === true,
+      });
     case "codex":
       // /me/status and the quota-cache refresh reach this fetch with no
       // caller-side pacing. Gate the start here so those paths, and any later
       // one, cannot burst the usage endpoint. Callers that already acquired the
       // shared gate wait at most one more interval.
       await throttleQuotaFetch();
-      return await getCodexUsage(accessToken, providerSpecificData);
+      return await getCodexUsage(accessToken, providerSpecificData, id || "");
     case "cursor":
-      return await getCursorUsage(accessToken || "", providerSpecificData);
+      return await getCursorUsage(accessToken || "", providerSpecificData, id || "");
     case "kiro":
     case "amazon-q":
       return await getKiroUsage(accessToken, providerSpecificData);
@@ -268,25 +270,25 @@ export async function getUsageForProvider(
     case "volcengine-coding-plan":
       return await getVolcenginePlanUsage(apiKey || "", provider, providerSpecificData);
     case "command-code":
-      return await getCommandCodeUsage(apiKey || accessToken || "");
+      return await getCommandCodeUsage(apiKey || accessToken || "", id || "");
     case "conol-web":
     case "cnl":
-      return await getConolUsage(apiKey || accessToken, providerSpecificData);
+      return await getConolUsage(apiKey || accessToken, providerSpecificData, id || "");
     case "syntx":
     case "stx":
-      return await getSyntxUsage(apiKey || accessToken, providerSpecificData);
+      return await getSyntxUsage(apiKey, accessToken, providerSpecificData, id || "");
     case "agentrouter":
       return await getAgentrouterUsage(id, connection);
     case "kilocode":
       return await getKilocodeUsage(id, connection);
     case "devin-cli":
       // Devin CLI tokens live in `accessToken` (oauth import) or `apiKey`.
-      return await getDevinCliUsage(apiKey || accessToken);
+      return await getDevinCliUsage(apiKey || accessToken, id || "");
     case "anthropic":
       return await getAnthropicApiKeyUsage(apiKey);
     case "chatplayground":
     case "cpl":
-      return await getChatPlaygroundUsage(apiKey || accessToken, providerSpecificData);
+      return await getChatPlaygroundUsage(apiKey || accessToken, providerSpecificData, id || "");
     default:
       return { message: `Usage API not implemented for ${provider}` };
   }

open-sse/services/usage/chatplayground.ts

@@ -5,6 +5,7 @@
  * Lifetime subscription provides 300 credits daily.
  */
 import { type UsageQuota } from "./quota.ts";
+import { fetchWithConnectionProxy } from "../connectionProxyFetch.ts";
 import { resolveChatPlaygroundAuth } from "../chatplaygroundAuth.ts";
 import { CHATPLAYGROUND_USER_URL } from "../chatplaygroundModels.ts";
 import { sanitizeErrorMessage } from "../../utils/error.ts";
@@ -234,7 +235,8 @@ export function buildChatPlaygroundQuotas(
 
 export async function getChatPlaygroundUsage(
   apiKey?: string,
-  providerSpecificData?: Record<string, unknown> | null
+  providerSpecificData?: Record<string, unknown> | null,
+  connectionId?: string
 ): Promise<{
   plan?: string;
   quotas?: Record<string, UsageQuota>;
@@ -246,10 +248,18 @@ export async function getChatPlaygroundUsage(
       providerSpecificData,
     });
 
-    const res = await fetch(CHATPLAYGROUND_USER_URL, {
-      method: "GET",
-      headers: auth.headers,
-    });
+    const res = connectionId
+      ? await fetchWithConnectionProxy(connectionId, CHATPLAYGROUND_USER_URL, {
+          method: "GET",
+          headers: auth.headers,
+        })
+      : await globalThis.fetch(CHATPLAYGROUND_USER_URL, {
+          method: "GET",
+          headers: auth.headers,
+        });
+    if (!res) {
+      throw new Error("proxied request failed");
+    }
 
     if (!res.ok) {
       const errText = await res.text().catch(() => "");