@tokenbuddy/tb-admin 1.0.14 → 1.0.27

package/src/upstream-balance-probe.ts

@@ -0,0 +1,505 @@
+export type BalanceCurrency = "USD" | "CNY" | null;
+export type BalanceSource =
+  | "deepseek"
+  | "stepfun"
+  | "siliconflow"
+  | "openrouter"
+  | "novita"
+  | "newapi_generic"
+  | "usage_generic"
+  | "unknown";
+
+export type BalanceProbeTemplate =
+  | "auto"
+  | "deepseek"
+  | "stepfun"
+  | "siliconflow"
+  | "openrouter"
+  | "novita"
+  | "newapi_generic"
+  | "usage_generic"
+  | "none";
+
+export interface BalanceSnapshot {
+  rawAmount: number | null;
+  amountUsdMicros: number | null;
+  currency: BalanceCurrency;
+  source: BalanceSource;
+  fetchedAt: number;
+  error?: {
+    httpStatus: number;
+    message: string;
+  };
+}
+
+export interface BalanceProbeConfig {
+  upstreamUrl?: string;
+  upstreamBalanceUrl?: string;
+  upstreamApiKey?: string;
+  upstreamUserId?: string;
+  upstreamBalanceProbe?: {
+    template?: BalanceProbeTemplate;
+    url?: string;
+    userId?: string;
+    rechargeUrl?: string;
+  };
+}
+
+export interface BalanceProbeOptions {
+  fetch?: typeof fetch;
+  now?: () => number;
+  timeoutMs?: number;
+  cnyUsdRate?: number;
+  cache?: BalanceProbeCache;
+}
+
+interface BalanceRequestPlan {
+  source: BalanceSource;
+  url: string;
+  currency: Exclude<BalanceCurrency, null>;
+  requiresUserId: boolean;
+}
+
+const DEFAULT_TIMEOUT_MS = 7000;
+const FAILURE_CACHE_TTL_MS = 30000;
+const DEFAULT_CNY_USD_RATE = 0.14;
+
+export class BalanceProbeCache {
+  private readonly entries = new Map<string, { expiresAt: number; snapshot: BalanceSnapshot }>();
+
+  public get(key: string, now: number): BalanceSnapshot | undefined {
+    const entry = this.entries.get(key);
+    if (!entry || entry.expiresAt <= now) {
+      if (entry) {
+        this.entries.delete(key);
+      }
+      return undefined;
+    }
+    return entry.snapshot;
+  }
+
+  public setFailure(key: string, snapshot: BalanceSnapshot, now: number): void {
+    this.entries.set(key, {
+      snapshot,
+      expiresAt: now + FAILURE_CACHE_TTL_MS
+    });
+  }
+}
+
+export const defaultBalanceProbeCache = new BalanceProbeCache();
+
+export async function probeUpstreamBalance(
+  config: BalanceProbeConfig,
+  options: BalanceProbeOptions = {}
+): Promise<BalanceSnapshot> {
+  const now = options.now || Date.now;
+  const fetchedAt = now();
+  const key = cacheKey(config);
+  const cache = options.cache || defaultBalanceProbeCache;
+  const cached = cache.get(key, fetchedAt);
+  if (cached) {
+    return cached;
+  }
+
+  const plan = requestPlan(config);
+  if (!plan) {
+    const snapshot = failure("unknown", fetchedAt, 0, "unsupported upstream - configure a balance parser");
+    cache.setFailure(key, snapshot, fetchedAt);
+    return snapshot;
+  }
+  if (!config.upstreamApiKey) {
+    const snapshot = failure(plan.source, fetchedAt, 0, "missing upstreamApiKey for balance probe");
+    cache.setFailure(key, snapshot, fetchedAt);
+    return snapshot;
+  }
+  const upstreamUserId = stringValue(config.upstreamBalanceProbe?.userId) || stringValue(config.upstreamUserId);
+  if (plan.requiresUserId && !upstreamUserId) {
+    const snapshot = failure(plan.source, fetchedAt, 0, "missing upstreamUserId for newapi upstream");
+    cache.setFailure(key, snapshot, fetchedAt);
+    return snapshot;
+  }
+
+  try {
+    const response = await fetchWithTimeout(
+      plan.url,
+      {
+        method: "GET",
+        headers: requestHeaders(config.upstreamApiKey, plan.requiresUserId ? upstreamUserId : undefined)
+      },
+      options.fetch || fetch,
+      options.timeoutMs ?? DEFAULT_TIMEOUT_MS
+    );
+    if (!response.ok) {
+      const snapshot = failure(plan.source, fetchedAt, response.status, httpErrorMessage(response.status));
+      cache.setFailure(key, snapshot, fetchedAt);
+      return snapshot;
+    }
+    const payload = await response.json() as unknown;
+    const snapshot = parseBalancePayload(payload, plan, fetchedAt, cnyUsdRate(options.cnyUsdRate));
+    if (snapshot.error) {
+      cache.setFailure(key, snapshot, fetchedAt);
+    }
+    return snapshot;
+  } catch (err) {
+    const snapshot = failure(plan.source, fetchedAt, 0, `network: ${err instanceof Error ? err.message : String(err)}`);
+    cache.setFailure(key, snapshot, fetchedAt);
+    return snapshot;
+  }
+}
+
+function requestPlan(config: BalanceProbeConfig): BalanceRequestPlan | undefined {
+  const template = balanceProbeTemplateValue(config.upstreamBalanceProbe?.template);
+  if (template === "none") {
+    return undefined;
+  }
+  const balanceUrl = stringValue(config.upstreamBalanceProbe?.url) || stringValue(config.upstreamBalanceUrl);
+  const upstreamUrl = stringValue(config.upstreamUrl);
+  if (template && template !== "auto") {
+    return explicitRequestPlan(template, balanceUrl, upstreamUrl);
+  }
+  const host = hostName(balanceUrl || upstreamUrl);
+  if (host === "api.deepseek.com") {
+    return { source: "deepseek", url: "https://api.deepseek.com/user/balance", currency: "CNY", requiresUserId: false };
+  }
+  if (host === "api.stepfun.ai" || host === "api.stepfun.com") {
+    return { source: "stepfun", url: "https://api.stepfun.com/v1/accounts", currency: "CNY", requiresUserId: false };
+  }
+  if (host === "api.siliconflow.cn" || host === "api.siliconflow.com") {
+    const currency = host.endsWith(".cn") ? "CNY" : "USD";
+    const url = host.endsWith(".cn") ? "https://api.siliconflow.cn/v1/user/info" : "https://api.siliconflow.com/v1/user/info";
+    return { source: "siliconflow", url, currency, requiresUserId: false };
+  }
+  if (host === "openrouter.ai") {
+    return { source: "openrouter", url: "https://openrouter.ai/api/v1/credits", currency: "USD", requiresUserId: false };
+  }
+  if (host === "api.novita.ai") {
+    return { source: "novita", url: "https://api.novita.ai/v3/user/balance", currency: "USD", requiresUserId: false };
+  }
+  if (balanceUrl && isUsageEndpoint(balanceUrl)) {
+    return { source: "usage_generic", url: balanceUrl, currency: "USD", requiresUserId: false };
+  }
+  if (balanceUrl) {
+    return { source: "newapi_generic", url: balanceUrl, currency: "USD", requiresUserId: true };
+  }
+  const genericUsageUrl = upstreamUrl ? usageUrl(upstreamUrl) : undefined;
+  if (genericUsageUrl) {
+    return { source: "usage_generic", url: genericUsageUrl, currency: "USD", requiresUserId: false };
+  }
+  return undefined;
+}
+
+function explicitRequestPlan(
+  template: Exclude<BalanceProbeTemplate, "auto" | "none">,
+  balanceUrl: string | undefined,
+  upstreamUrl: string | undefined
+): BalanceRequestPlan | undefined {
+  if (template === "deepseek") {
+    return { source: "deepseek", url: balanceUrl || "https://api.deepseek.com/user/balance", currency: "CNY", requiresUserId: false };
+  }
+  if (template === "stepfun") {
+    return { source: "stepfun", url: balanceUrl || "https://api.stepfun.com/v1/accounts", currency: "CNY", requiresUserId: false };
+  }
+  if (template === "siliconflow") {
+    const host = hostName(balanceUrl || upstreamUrl);
+    const currency = host.endsWith(".com") ? "USD" : "CNY";
+    const url = balanceUrl || (currency === "USD" ? "https://api.siliconflow.com/v1/user/info" : "https://api.siliconflow.cn/v1/user/info");
+    return { source: "siliconflow", url, currency, requiresUserId: false };
+  }
+  if (template === "openrouter") {
+    return { source: "openrouter", url: balanceUrl || "https://openrouter.ai/api/v1/credits", currency: "USD", requiresUserId: false };
+  }
+  if (template === "novita") {
+    return { source: "novita", url: balanceUrl || "https://api.novita.ai/v3/user/balance", currency: "USD", requiresUserId: false };
+  }
+  if (template === "newapi_generic") {
+    return balanceUrl ? { source: "newapi_generic", url: balanceUrl, currency: "USD", requiresUserId: true } : undefined;
+  }
+  if (template === "usage_generic") {
+    const url = balanceUrl || (upstreamUrl ? usageUrl(upstreamUrl) : undefined);
+    return url ? { source: "usage_generic", url, currency: "USD", requiresUserId: false } : undefined;
+  }
+  return undefined;
+}
+
+function requestHeaders(upstreamApiKey: string, upstreamUserId: string | undefined): Record<string, string> {
+  const headers: Record<string, string> = {
+    "Authorization": `Bearer ${upstreamApiKey}`
+  };
+  if (upstreamUserId) {
+    headers["New-Api-User"] = upstreamUserId;
+  }
+  return headers;
+}
+
+async function fetchWithTimeout(
+  url: string,
+  init: RequestInit,
+  fetchFn: typeof fetch,
+  timeoutMs: number
+): Promise<Response> {
+  const controller = new AbortController();
+  const timer = setTimeout(() => controller.abort(), timeoutMs);
+  try {
+    return await fetchFn(url, { ...init, signal: controller.signal });
+  } finally {
+    clearTimeout(timer);
+  }
+}
+
+function parseBalancePayload(
+  payload: unknown,
+  plan: BalanceRequestPlan,
+  fetchedAt: number,
+  cnyRate: number
+): BalanceSnapshot {
+  const object = objectValue(payload);
+  if (!object) {
+    return failure(plan.source, fetchedAt, 200, "balance response must be an object");
+  }
+  if (plan.source === "deepseek") {
+    return parseDeepSeek(object, fetchedAt, cnyRate);
+  }
+  if (plan.source === "stepfun") {
+    return amountSnapshot(numberFrom(object.balance), "CNY", "stepfun", fetchedAt, cnyRate);
+  }
+  if (plan.source === "siliconflow") {
+    return parseSiliconFlow(object, plan.currency, fetchedAt, cnyRate);
+  }
+  if (plan.source === "openrouter") {
+    return parseOpenRouter(object, fetchedAt, cnyRate);
+  }
+  if (plan.source === "novita") {
+    const availableBalance = numberFrom(object.availableBalance);
+    return amountSnapshot(availableBalance === null ? null : availableBalance / 10000, "USD", "novita", fetchedAt, cnyRate);
+  }
+  if (plan.source === "newapi_generic") {
+    return parseNewApi(object, fetchedAt, cnyRate);
+  }
+  if (plan.source === "usage_generic") {
+    return parseUsageGeneric(object, fetchedAt, cnyRate);
+  }
+  return failure("unknown", fetchedAt, 0, "unsupported upstream - configure a balance parser");
+}
+
+function parseDeepSeek(payload: Record<string, unknown>, fetchedAt: number, cnyRate: number): BalanceSnapshot {
+  const infos = Array.isArray(payload.balance_infos) ? payload.balance_infos : [];
+  const first = objectValue(infos[0]);
+  const currency = currencyValue(first?.currency) || "CNY";
+  const rawAmount = numberFrom(first?.total_balance);
+  if (payload.is_available === false) {
+    return {
+      ...amountSnapshot(rawAmount, currency, "deepseek", fetchedAt, cnyRate),
+      error: { httpStatus: 200, message: "Insufficient balance" }
+    };
+  }
+  if (!first) {
+    return amountSnapshot(null, currency, "deepseek", fetchedAt, cnyRate);
+  }
+  if (rawAmount === null) {
+    return failure("deepseek", fetchedAt, 200, "missing field: balance_infos[].total_balance", currency);
+  }
+  return amountSnapshot(rawAmount, currency, "deepseek", fetchedAt, cnyRate);
+}
+
+function parseSiliconFlow(
+  payload: Record<string, unknown>,
+  currency: Exclude<BalanceCurrency, null>,
+  fetchedAt: number,
+  cnyRate: number
+): BalanceSnapshot {
+  const code = numberFrom(payload.code);
+  if (code !== null && code !== 20000) {
+    return failure("siliconflow", fetchedAt, 200, `upstream code: ${code}`, currency);
+  }
+  const data = objectValue(payload.data);
+  if (!data) {
+    return failure("siliconflow", fetchedAt, 200, "missing field: data", currency);
+  }
+  const snapshot = amountSnapshot(numberFrom(data.totalBalance), currency, "siliconflow", fetchedAt, cnyRate);
+  const status = stringValue(data.status);
+  if (status && status !== "ok") {
+    return {
+      ...snapshot,
+      error: { httpStatus: 200, message: `upstream status: ${status}` }
+    };
+  }
+  return snapshot;
+}
+
+function parseOpenRouter(payload: Record<string, unknown>, fetchedAt: number, cnyRate: number): BalanceSnapshot {
+  const data = objectValue(payload.data) || payload;
+  const totalCredits = numberFrom(data.total_credits);
+  const totalUsage = numberFrom(data.total_usage);
+  if (totalCredits === null || totalUsage === null) {
+    return failure("openrouter", fetchedAt, 200, "missing field: data.total_credits or data.total_usage", "USD");
+  }
+  return amountSnapshot(Math.max(0, totalCredits - totalUsage), "USD", "openrouter", fetchedAt, cnyRate);
+}
+
+function parseNewApi(payload: Record<string, unknown>, fetchedAt: number, cnyRate: number): BalanceSnapshot {
+  const data = objectValue(payload.data);
+  if (!data) {
+    return failure("newapi_generic", fetchedAt, 200, "missing field: data", "USD");
+  }
+  const quota = numberFrom(data.quota);
+  const usedQuota = numberFrom(data.used_quota);
+  if (quota === null || usedQuota === null) {
+    return failure("newapi_generic", fetchedAt, 200, "missing field: data.quota or data.used_quota", "USD");
+  }
+  return amountSnapshot((quota - usedQuota) / 500000, "USD", "newapi_generic", fetchedAt, cnyRate);
+}
+
+function parseUsageGeneric(payload: Record<string, unknown>, fetchedAt: number, cnyRate: number): BalanceSnapshot {
+  const quota = objectValue(payload.quota);
+  const rawAmount = numberFrom(payload.remaining) ?? numberFrom(quota?.remaining) ?? numberFrom(payload.balance);
+  const currency = currencyValue(payload.unit) || currencyValue(quota?.unit) || "USD";
+  if (rawAmount === null) {
+    return failure("usage_generic", fetchedAt, 200, "missing field: remaining, quota.remaining, or balance", currency);
+  }
+  const snapshot = amountSnapshot(rawAmount, currency, "usage_generic", fetchedAt, cnyRate);
+  const isValid = booleanValue(payload.is_active) ?? booleanValue(payload.isValid) ?? true;
+  if (!isValid) {
+    return {
+      ...snapshot,
+      error: { httpStatus: 200, message: "upstream key is not active" }
+    };
+  }
+  return snapshot;
+}
+
+function amountSnapshot(
+  rawAmount: number | null,
+  currency: Exclude<BalanceCurrency, null>,
+  source: BalanceSource,
+  fetchedAt: number,
+  cnyRate: number
+): BalanceSnapshot {
+  return {
+    rawAmount,
+    amountUsdMicros: rawAmount === null ? null : Math.round(rawAmount * (currency === "CNY" ? cnyRate : 1) * 1000000),
+    currency,
+    source,
+    fetchedAt
+  };
+}
+
+function failure(
+  source: BalanceSource,
+  fetchedAt: number,
+  httpStatus: number,
+  message: string,
+  currency: BalanceCurrency = null
+): BalanceSnapshot {
+  return {
+    rawAmount: null,
+    amountUsdMicros: null,
+    currency,
+    source,
+    fetchedAt,
+    error: { httpStatus, message }
+  };
+}
+
+function httpErrorMessage(status: number): string {
+  if (status === 401) {
+    return "unauthorized: check upstreamApiKey";
+  }
+  if (status === 403) {
+    return "forbidden: upstream rejected the credentials";
+  }
+  if (status === 429) {
+    return "rate limited";
+  }
+  if (status >= 500) {
+    return "upstream 5xx";
+  }
+  return `upstream http ${status}`;
+}
+
+function cnyUsdRate(value: number | undefined): number {
+  if (value !== undefined && Number.isFinite(value) && value > 0) {
+    return value;
+  }
+  const envValue = Number(process.env.TB_CNY_USD_RATE);
+  return Number.isFinite(envValue) && envValue > 0 ? envValue : DEFAULT_CNY_USD_RATE;
+}
+
+function cacheKey(config: BalanceProbeConfig): string {
+  return [
+    stringValue(config.upstreamBalanceProbe?.template),
+    stringValue(config.upstreamBalanceProbe?.url),
+    stringValue(config.upstreamBalanceProbe?.userId),
+    stringValue(config.upstreamBalanceUrl),
+    stringValue(config.upstreamUrl),
+    stringValue(config.upstreamApiKey),
+    stringValue(config.upstreamUserId)
+  ].join("|");
+}
+
+function balanceProbeTemplateValue(value: unknown): BalanceProbeTemplate | undefined {
+  return value === "auto" ||
+    value === "deepseek" ||
+    value === "stepfun" ||
+    value === "siliconflow" ||
+    value === "openrouter" ||
+    value === "novita" ||
+    value === "newapi_generic" ||
+    value === "usage_generic" ||
+    value === "none"
+    ? value
+    : undefined;
+}
+
+function hostName(value: string | undefined): string {
+  if (!value) {
+    return "";
+  }
+  try {
+    return new URL(value).hostname.replace(/^www\./, "");
+  } catch {
+    return "";
+  }
+}
+
+function usageUrl(value: string): string | undefined {
+  try {
+    const url = new URL(value);
+    url.pathname = "/v1/usage";
+    url.search = "";
+    url.hash = "";
+    return url.toString();
+  } catch {
+    return undefined;
+  }
+}
+
+function isUsageEndpoint(value: string): boolean {
+  try {
+    return new URL(value).pathname.replace(/\/+$/, "") === "/v1/usage";
+  } catch {
+    return false;
+  }
+}
+
+function objectValue(value: unknown): Record<string, unknown> | undefined {
+  return value && typeof value === "object" && !Array.isArray(value)
+    ? value as Record<string, unknown>
+    : undefined;
+}
+
+function stringValue(value: unknown): string | undefined {
+  return typeof value === "string" && value.trim() ? value.trim() : undefined;
+}
+
+function numberFrom(value: unknown): number | null {
+  const parsed = typeof value === "number" ? value : typeof value === "string" ? Number(value) : Number.NaN;
+  return Number.isFinite(parsed) ? parsed : null;
+}
+
+function booleanValue(value: unknown): boolean | undefined {
+  return typeof value === "boolean" ? value : undefined;
+}
+
+function currencyValue(value: unknown): Exclude<BalanceCurrency, null> | undefined {
+  return value === "USD" || value === "CNY" ? value : undefined;
+}