From 780f8c26686c9daacb5aa038035d5491af6f55e4 Mon Sep 17 00:00:00 2001 From: jonisjongithub <86072337+jonisjongithub@users.noreply.github.com> Date: Tue, 27 Jan 2026 23:51:20 -0800 Subject: [PATCH] feat: add API key spending limit/rate limit tracking - Change default low threshold to 5 DIEM (critical stays at 2) - Add rate limit tracking from x-ratelimit-* headers - Track remaining requests/tokens vs limits - Add configurable lowRateLimitPercent (default 10%) and criticalRateLimitPercent (default 5%) - Generate warnings when approaching API key spending caps - Show rate limit info in /status output - Add comprehensive tests for rate limit scenarios Co-authored-by: jonisjongithub Co-authored-by: Clawdbot --- src/config/zod-schema.core.ts | 6 +- src/providers/venice-balance.test.ts | 138 ++++++++++++++++++-- src/providers/venice-balance.ts | 182 +++++++++++++++++++++++---- 3 files changed, 295 insertions(+), 31 deletions(-) diff --git a/src/config/zod-schema.core.ts b/src/config/zod-schema.core.ts index 3fcc95fb0..06e920c2b 100644 --- a/src/config/zod-schema.core.ts +++ b/src/config/zod-schema.core.ts @@ -76,10 +76,14 @@ export const VeniceBalanceWarningSchema = z .object({ /** Enable balance warnings (default: true). */ enabled: z.boolean().optional(), - /** Warn when DIEM balance falls below this threshold (default: 10). */ + /** Warn when DIEM balance falls below this threshold (default: 5). */ lowDiemThreshold: z.number().nonnegative().optional(), /** Critical warning when DIEM balance falls below this threshold (default: 2). */ criticalDiemThreshold: z.number().nonnegative().optional(), + /** Warn when remaining requests falls below this percentage of limit (default: 10). */ + lowRateLimitPercent: z.number().min(0).max(100).optional(), + /** Critical warning when remaining requests falls below this percentage (default: 5). */ + criticalRateLimitPercent: z.number().min(0).max(100).optional(), /** Show balance in /status output (default: true). */ showInStatus: z.boolean().optional(), }) diff --git a/src/providers/venice-balance.test.ts b/src/providers/venice-balance.test.ts index 9f205a1c8..6aa76ba61 100644 --- a/src/providers/venice-balance.test.ts +++ b/src/providers/venice-balance.test.ts @@ -5,6 +5,7 @@ import { getVeniceBalance, clearVeniceBalance, evaluateBalanceStatus, + evaluateBalanceDetailed, formatVeniceBalanceStatus, generateBalanceWarning, isVeniceBalanceError, @@ -69,6 +70,26 @@ describe("venice-balance", () => { expect(balance).toBeNull(); }); + + it("extracts rate limit headers", () => { + const headers = new Headers(); + headers.set("x-venice-balance-diem", "44.32"); + headers.set("x-ratelimit-limit-requests", "500"); + headers.set("x-ratelimit-remaining-requests", "499"); + headers.set("x-ratelimit-limit-tokens", "5000000"); + headers.set("x-ratelimit-remaining-tokens", "4999000"); + headers.set("x-ratelimit-reset-requests", "1769585280000"); + + const balance = extractVeniceBalance(headers); + + expect(balance).not.toBeNull(); + expect(balance?.rateLimit).toBeDefined(); + expect(balance?.rateLimit?.limitRequests).toBe(500); + expect(balance?.rateLimit?.remainingRequests).toBe(499); + expect(balance?.rateLimit?.limitTokens).toBe(5000000); + expect(balance?.rateLimit?.remainingTokens).toBe(4999000); + expect(balance?.rateLimit?.resetAt).toBe(1769585280000); + }); }); describe("updateVeniceBalance / getVeniceBalance", () => { @@ -100,12 +121,12 @@ describe("venice-balance", () => { expect(evaluateBalanceStatus(balance)).toBe("ok"); }); - it("returns 'low' for balance below low threshold", () => { - const balance: VeniceBalance = { diem: 5, lastChecked: Date.now() }; + it("returns 'low' for balance below low threshold (5 DIEM)", () => { + const balance: VeniceBalance = { diem: 4, lastChecked: Date.now() }; expect(evaluateBalanceStatus(balance)).toBe("low"); }); - it("returns 'critical' for balance below critical threshold", () => { + it("returns 'critical' for balance below critical threshold (2 DIEM)", () => { const balance: VeniceBalance = { diem: 1, lastChecked: Date.now() }; expect(evaluateBalanceStatus(balance)).toBe("critical"); }); @@ -129,6 +150,77 @@ describe("venice-balance", () => { }); }); + describe("evaluateBalanceDetailed - rate limits", () => { + it("returns critical when request limit is nearly exhausted", () => { + const balance: VeniceBalance = { + diem: 50, + lastChecked: Date.now(), + rateLimit: { + limitRequests: 500, + remainingRequests: 10, // 2% remaining, below 5% critical threshold + }, + }; + const evaluation = evaluateBalanceDetailed(balance); + expect(evaluation.status).toBe("critical"); + expect(evaluation.reason).toBe("rate_limit"); + }); + + it("returns low when request limit is getting low", () => { + const balance: VeniceBalance = { + diem: 50, + lastChecked: Date.now(), + rateLimit: { + limitRequests: 500, + remainingRequests: 40, // 8% remaining, below 10% low threshold + }, + }; + const evaluation = evaluateBalanceDetailed(balance); + expect(evaluation.status).toBe("low"); + expect(evaluation.reason).toBe("rate_limit"); + }); + + it("returns depleted when request limit is exhausted", () => { + const balance: VeniceBalance = { + diem: 50, + lastChecked: Date.now(), + rateLimit: { + limitRequests: 500, + remainingRequests: 0, + }, + }; + const evaluation = evaluateBalanceDetailed(balance); + expect(evaluation.status).toBe("depleted"); + expect(evaluation.reason).toBe("rate_limit"); + }); + + it("DIEM takes priority over rate limit for depleted status", () => { + const balance: VeniceBalance = { + diem: 0, + lastChecked: Date.now(), + rateLimit: { + limitRequests: 500, + remainingRequests: 500, // plenty of requests + }, + }; + const evaluation = evaluateBalanceDetailed(balance); + expect(evaluation.status).toBe("depleted"); + expect(evaluation.reason).toBe("diem"); + }); + + it("returns ok when both DIEM and rate limits are healthy", () => { + const balance: VeniceBalance = { + diem: 50, + lastChecked: Date.now(), + rateLimit: { + limitRequests: 500, + remainingRequests: 400, // 80% remaining + }, + }; + const evaluation = evaluateBalanceDetailed(balance); + expect(evaluation.status).toBe("ok"); + }); + }); + describe("formatVeniceBalanceStatus", () => { it("returns formatted status for OK balance", () => { const balance: VeniceBalance = { diem: 50, lastChecked: Date.now() }; @@ -163,16 +255,16 @@ describe("venice-balance", () => { expect(generateBalanceWarning(balance)).toBeNull(); }); - it("returns warning for low balance", () => { - const balance: VeniceBalance = { diem: 5, lastChecked: Date.now() }; + it("returns warning for low DIEM balance", () => { + const balance: VeniceBalance = { diem: 4, lastChecked: Date.now() }; const warning = generateBalanceWarning(balance); expect(warning).toContain("⚠️"); expect(warning).toContain("low"); - expect(warning).toContain("5.00"); + expect(warning).toContain("4.00"); }); - it("returns critical warning for critical balance", () => { + it("returns critical warning for critical DIEM balance", () => { const balance: VeniceBalance = { diem: 1, lastChecked: Date.now() }; const warning = generateBalanceWarning(balance); @@ -180,7 +272,7 @@ describe("venice-balance", () => { expect(warning).toContain("critical"); }); - it("returns depleted warning for zero balance", () => { + it("returns depleted warning for zero DIEM balance", () => { const balance: VeniceBalance = { diem: 0, lastChecked: Date.now() }; const warning = generateBalanceWarning(balance); @@ -194,6 +286,36 @@ describe("venice-balance", () => { expect(generateBalanceWarning(balance, thresholds)).toBeNull(); }); + + it("returns warning for low rate limit", () => { + const balance: VeniceBalance = { + diem: 50, + lastChecked: Date.now(), + rateLimit: { + limitRequests: 500, + remainingRequests: 40, // 8% remaining + }, + }; + const warning = generateBalanceWarning(balance); + + expect(warning).toContain("⚠️"); + expect(warning).toContain("API key limit"); + }); + + it("returns critical warning for exhausted rate limit", () => { + const balance: VeniceBalance = { + diem: 50, + lastChecked: Date.now(), + rateLimit: { + limitRequests: 500, + remainingRequests: 0, + }, + }; + const warning = generateBalanceWarning(balance); + + expect(warning).toContain("❌"); + expect(warning).toContain("limit reached"); + }); }); describe("isVeniceBalanceError", () => { diff --git a/src/providers/venice-balance.ts b/src/providers/venice-balance.ts index 74af346fb..7d9c7a1b2 100644 --- a/src/providers/venice-balance.ts +++ b/src/providers/venice-balance.ts @@ -20,19 +20,40 @@ export interface VeniceBalance { lastChecked: number; /** Provider identifier (for multi-key scenarios) */ providerId?: string; + /** Rate limit info (tracks API key spending cap usage) */ + rateLimit?: { + /** Max requests allowed in current window */ + limitRequests?: number; + /** Remaining requests in current window */ + remainingRequests?: number; + /** Max tokens allowed in current window */ + limitTokens?: number; + /** Remaining tokens in current window */ + remainingTokens?: number; + /** Reset timestamp (unix ms) */ + resetAt?: number; + }; } export interface VeniceBalanceThresholds { enabled: boolean; + /** Warn when DIEM balance falls below this value (default: 5) */ lowDiemThreshold: number; + /** Critical warning when DIEM balance falls below this value (default: 2) */ criticalDiemThreshold: number; + /** Warn when remaining requests falls below this percentage of limit (default: 10%) */ + lowRateLimitPercent: number; + /** Critical warning when remaining requests falls below this percentage (default: 5%) */ + criticalRateLimitPercent: number; showInStatus: boolean; } export const DEFAULT_VENICE_BALANCE_THRESHOLDS: VeniceBalanceThresholds = { enabled: true, - lowDiemThreshold: 10, + lowDiemThreshold: 5, criticalDiemThreshold: 2, + lowRateLimitPercent: 10, + criticalRateLimitPercent: 5, showInStatus: true, }; @@ -63,10 +84,28 @@ export function extractVeniceBalance(headers: Headers | Record): const usdRaw = get("x-venice-balance-usd"); const vcuRaw = get("x-venice-balance-vcu"); + // Extract rate limit headers (tracks API key spending cap) + const limitRequestsRaw = get("x-ratelimit-limit-requests"); + const remainingRequestsRaw = get("x-ratelimit-remaining-requests"); + const limitTokensRaw = get("x-ratelimit-limit-tokens"); + const remainingTokensRaw = get("x-ratelimit-remaining-tokens"); + const resetRequestsRaw = get("x-ratelimit-reset-requests"); + + const rateLimit = (limitRequestsRaw || remainingRequestsRaw || limitTokensRaw || remainingTokensRaw) + ? { + limitRequests: limitRequestsRaw ? parseInt(limitRequestsRaw, 10) : undefined, + remainingRequests: remainingRequestsRaw ? parseInt(remainingRequestsRaw, 10) : undefined, + limitTokens: limitTokensRaw ? parseInt(limitTokensRaw, 10) : undefined, + remainingTokens: remainingTokensRaw ? parseInt(remainingTokensRaw, 10) : undefined, + resetAt: resetRequestsRaw ? parseInt(resetRequestsRaw, 10) : undefined, + } + : undefined; + return { diem, usd: usdRaw ? parseFloat(usdRaw) : undefined, vcu: vcuRaw ? parseFloat(vcuRaw) : undefined, + rateLimit, lastChecked: Date.now(), }; } @@ -119,15 +158,82 @@ export function onVeniceBalanceUpdate( */ export type BalanceStatus = "ok" | "low" | "critical" | "depleted" | "unknown"; +export interface BalanceEvaluation { + status: BalanceStatus; + /** Why this status was assigned */ + reason: "diem" | "rate_limit" | "unknown"; + /** Human-readable details */ + details?: string; +} + export function evaluateBalanceStatus( balance: VeniceBalance | null, thresholds: VeniceBalanceThresholds = DEFAULT_VENICE_BALANCE_THRESHOLDS, ): BalanceStatus { - if (!balance || balance.diem === undefined) return "unknown"; - if (balance.diem <= 0) return "depleted"; - if (balance.diem < thresholds.criticalDiemThreshold) return "critical"; - if (balance.diem < thresholds.lowDiemThreshold) return "low"; - return "ok"; + return evaluateBalanceDetailed(balance, thresholds).status; +} + +/** + * Evaluate balance with detailed reason. + */ +export function evaluateBalanceDetailed( + balance: VeniceBalance | null, + thresholds: VeniceBalanceThresholds = DEFAULT_VENICE_BALANCE_THRESHOLDS, +): BalanceEvaluation { + if (!balance) return { status: "unknown", reason: "unknown" }; + + // Check DIEM balance first (most common) + if (balance.diem !== undefined) { + if (balance.diem <= 0) { + return { status: "depleted", reason: "diem", details: "DIEM balance is zero" }; + } + if (balance.diem < thresholds.criticalDiemThreshold) { + return { status: "critical", reason: "diem", details: `DIEM below ${thresholds.criticalDiemThreshold}` }; + } + if (balance.diem < thresholds.lowDiemThreshold) { + return { status: "low", reason: "diem", details: `DIEM below ${thresholds.lowDiemThreshold}` }; + } + } + + // Check rate limits (API key spending cap) + if (balance.rateLimit) { + const { limitRequests, remainingRequests, limitTokens, remainingTokens } = balance.rateLimit; + + // Check request limit + if (limitRequests !== undefined && remainingRequests !== undefined && limitRequests > 0) { + const remainingPercent = (remainingRequests / limitRequests) * 100; + if (remainingRequests <= 0) { + return { status: "depleted", reason: "rate_limit", details: "Request limit exhausted" }; + } + if (remainingPercent < thresholds.criticalRateLimitPercent) { + return { status: "critical", reason: "rate_limit", details: `Only ${remainingRequests}/${limitRequests} requests remaining` }; + } + if (remainingPercent < thresholds.lowRateLimitPercent) { + return { status: "low", reason: "rate_limit", details: `${remainingRequests}/${limitRequests} requests remaining` }; + } + } + + // Check token limit + if (limitTokens !== undefined && remainingTokens !== undefined && limitTokens > 0) { + const remainingPercent = (remainingTokens / limitTokens) * 100; + if (remainingTokens <= 0) { + return { status: "depleted", reason: "rate_limit", details: "Token limit exhausted" }; + } + if (remainingPercent < thresholds.criticalRateLimitPercent) { + return { status: "critical", reason: "rate_limit", details: `Only ${remainingTokens.toLocaleString()}/${limitTokens.toLocaleString()} tokens remaining` }; + } + if (remainingPercent < thresholds.lowRateLimitPercent) { + return { status: "low", reason: "rate_limit", details: `${remainingTokens.toLocaleString()}/${limitTokens.toLocaleString()} tokens remaining` }; + } + } + } + + // If we have DIEM info, we're OK + if (balance.diem !== undefined) { + return { status: "ok", reason: "diem" }; + } + + return { status: "unknown", reason: "unknown" }; } /** @@ -140,7 +246,7 @@ export function formatVeniceBalanceStatus( if (!thresholds.showInStatus) return null; if (!balance) return null; - const status = evaluateBalanceStatus(balance, thresholds); + const evaluation = evaluateBalanceDetailed(balance, thresholds); const diemLabel = balance.diem !== undefined ? balance.diem.toFixed(2) : "?"; const ageMs = Date.now() - balance.lastChecked; const ageLabel = formatAge(ageMs); @@ -155,14 +261,26 @@ export function formatVeniceBalanceStatus( const parts = [ `DIEM: ${diemLabel}`, - `Status: ${statusEmoji[status]} ${status.toUpperCase()}`, - `(${ageLabel})`, ]; if (balance.usd !== undefined) { - parts.splice(1, 0, `USD: $${balance.usd.toFixed(2)}`); + parts.push(`USD: $${balance.usd.toFixed(2)}`); } + // Add rate limit info if available + if (balance.rateLimit?.remainingRequests !== undefined && balance.rateLimit?.limitRequests !== undefined) { + parts.push(`Requests: ${balance.rateLimit.remainingRequests}/${balance.rateLimit.limitRequests}`); + } + + parts.push(`Status: ${statusEmoji[evaluation.status]} ${evaluation.status.toUpperCase()}`); + + // Add reason if not OK + if (evaluation.status !== "ok" && evaluation.details) { + parts.push(`(${evaluation.details})`); + } + + parts.push(`(${ageLabel})`); + return parts.join(" · "); } @@ -175,21 +293,39 @@ export function generateBalanceWarning( thresholds: VeniceBalanceThresholds = DEFAULT_VENICE_BALANCE_THRESHOLDS, ): string | null { if (!thresholds.enabled) return null; - if (!balance || balance.diem === undefined) return null; + if (!balance) return null; - const status = evaluateBalanceStatus(balance, thresholds); - const diemLabel = balance.diem.toFixed(2); + const evaluation = evaluateBalanceDetailed(balance, thresholds); + const diemLabel = balance.diem !== undefined ? balance.diem.toFixed(2) : "?"; - switch (status) { - case "critical": - return `🚨 Venice balance critical: ${diemLabel} DIEM remaining. Consider topping up at https://venice.ai/settings/billing`; - case "low": - return `⚠️ Venice balance low: ${diemLabel} DIEM remaining`; - case "depleted": - return `❌ Venice balance depleted (0.00 DIEM). Top up at https://venice.ai/settings/billing`; - default: - return null; + // DIEM-based warnings + if (evaluation.reason === "diem") { + switch (evaluation.status) { + case "critical": + return `🚨 Venice balance critical: ${diemLabel} DIEM remaining. Consider topping up at https://venice.ai/settings/billing`; + case "low": + return `⚠️ Venice balance low: ${diemLabel} DIEM remaining`; + case "depleted": + return `❌ Venice balance depleted (0.00 DIEM). Top up at https://venice.ai/settings/billing`; + } } + + // Rate limit / spending cap warnings + if (evaluation.reason === "rate_limit") { + const rl = balance.rateLimit; + const resetInfo = rl?.resetAt ? ` (resets ${formatAge(rl.resetAt - Date.now()).replace(" ago", "")})` : ""; + + switch (evaluation.status) { + case "critical": + return `🚨 Venice API key limit critical: ${evaluation.details}${resetInfo}. Consider increasing your spending cap.`; + case "low": + return `⚠️ Venice API key limit low: ${evaluation.details}${resetInfo}`; + case "depleted": + return `❌ Venice API key limit reached: ${evaluation.details}. Increase your spending cap at https://venice.ai/settings/api`; + } + } + + return null; } /** @@ -351,6 +487,8 @@ export function resolveVeniceBalanceThresholds( enabled: override.enabled ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.enabled, lowDiemThreshold: override.lowDiemThreshold ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.lowDiemThreshold, criticalDiemThreshold: override.criticalDiemThreshold ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.criticalDiemThreshold, + lowRateLimitPercent: override.lowRateLimitPercent ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.lowRateLimitPercent, + criticalRateLimitPercent: override.criticalRateLimitPercent ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.criticalRateLimitPercent, showInStatus: override.showInStatus ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.showInStatus, }; }