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 <jonisjongithub@users.noreply.github.com>
Co-authored-by: Clawdbot <bot@clawd.bot>
This commit is contained in:
jonisjongithub 2026-01-27 23:51:20 -08:00
parent 17770fb1c3
commit 780f8c2668
3 changed files with 295 additions and 31 deletions

View File

@ -76,10 +76,14 @@ export const VeniceBalanceWarningSchema = z
.object({ .object({
/** Enable balance warnings (default: true). */ /** Enable balance warnings (default: true). */
enabled: z.boolean().optional(), 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(), lowDiemThreshold: z.number().nonnegative().optional(),
/** Critical warning when DIEM balance falls below this threshold (default: 2). */ /** Critical warning when DIEM balance falls below this threshold (default: 2). */
criticalDiemThreshold: z.number().nonnegative().optional(), 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). */ /** Show balance in /status output (default: true). */
showInStatus: z.boolean().optional(), showInStatus: z.boolean().optional(),
}) })

View File

@ -5,6 +5,7 @@ import {
getVeniceBalance, getVeniceBalance,
clearVeniceBalance, clearVeniceBalance,
evaluateBalanceStatus, evaluateBalanceStatus,
evaluateBalanceDetailed,
formatVeniceBalanceStatus, formatVeniceBalanceStatus,
generateBalanceWarning, generateBalanceWarning,
isVeniceBalanceError, isVeniceBalanceError,
@ -69,6 +70,26 @@ describe("venice-balance", () => {
expect(balance).toBeNull(); 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", () => { describe("updateVeniceBalance / getVeniceBalance", () => {
@ -100,12 +121,12 @@ describe("venice-balance", () => {
expect(evaluateBalanceStatus(balance)).toBe("ok"); expect(evaluateBalanceStatus(balance)).toBe("ok");
}); });
it("returns 'low' for balance below low threshold", () => { it("returns 'low' for balance below low threshold (5 DIEM)", () => {
const balance: VeniceBalance = { diem: 5, lastChecked: Date.now() }; const balance: VeniceBalance = { diem: 4, lastChecked: Date.now() };
expect(evaluateBalanceStatus(balance)).toBe("low"); 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() }; const balance: VeniceBalance = { diem: 1, lastChecked: Date.now() };
expect(evaluateBalanceStatus(balance)).toBe("critical"); 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", () => { describe("formatVeniceBalanceStatus", () => {
it("returns formatted status for OK balance", () => { it("returns formatted status for OK balance", () => {
const balance: VeniceBalance = { diem: 50, lastChecked: Date.now() }; const balance: VeniceBalance = { diem: 50, lastChecked: Date.now() };
@ -163,16 +255,16 @@ describe("venice-balance", () => {
expect(generateBalanceWarning(balance)).toBeNull(); expect(generateBalanceWarning(balance)).toBeNull();
}); });
it("returns warning for low balance", () => { it("returns warning for low DIEM balance", () => {
const balance: VeniceBalance = { diem: 5, lastChecked: Date.now() }; const balance: VeniceBalance = { diem: 4, lastChecked: Date.now() };
const warning = generateBalanceWarning(balance); const warning = generateBalanceWarning(balance);
expect(warning).toContain("⚠️"); expect(warning).toContain("⚠️");
expect(warning).toContain("low"); 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 balance: VeniceBalance = { diem: 1, lastChecked: Date.now() };
const warning = generateBalanceWarning(balance); const warning = generateBalanceWarning(balance);
@ -180,7 +272,7 @@ describe("venice-balance", () => {
expect(warning).toContain("critical"); 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 balance: VeniceBalance = { diem: 0, lastChecked: Date.now() };
const warning = generateBalanceWarning(balance); const warning = generateBalanceWarning(balance);
@ -194,6 +286,36 @@ describe("venice-balance", () => {
expect(generateBalanceWarning(balance, thresholds)).toBeNull(); 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", () => { describe("isVeniceBalanceError", () => {

View File

@ -20,19 +20,40 @@ export interface VeniceBalance {
lastChecked: number; lastChecked: number;
/** Provider identifier (for multi-key scenarios) */ /** Provider identifier (for multi-key scenarios) */
providerId?: string; 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 { export interface VeniceBalanceThresholds {
enabled: boolean; enabled: boolean;
/** Warn when DIEM balance falls below this value (default: 5) */
lowDiemThreshold: number; lowDiemThreshold: number;
/** Critical warning when DIEM balance falls below this value (default: 2) */
criticalDiemThreshold: number; 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; showInStatus: boolean;
} }
export const DEFAULT_VENICE_BALANCE_THRESHOLDS: VeniceBalanceThresholds = { export const DEFAULT_VENICE_BALANCE_THRESHOLDS: VeniceBalanceThresholds = {
enabled: true, enabled: true,
lowDiemThreshold: 10, lowDiemThreshold: 5,
criticalDiemThreshold: 2, criticalDiemThreshold: 2,
lowRateLimitPercent: 10,
criticalRateLimitPercent: 5,
showInStatus: true, showInStatus: true,
}; };
@ -63,10 +84,28 @@ export function extractVeniceBalance(headers: Headers | Record<string, string>):
const usdRaw = get("x-venice-balance-usd"); const usdRaw = get("x-venice-balance-usd");
const vcuRaw = get("x-venice-balance-vcu"); 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 { return {
diem, diem,
usd: usdRaw ? parseFloat(usdRaw) : undefined, usd: usdRaw ? parseFloat(usdRaw) : undefined,
vcu: vcuRaw ? parseFloat(vcuRaw) : undefined, vcu: vcuRaw ? parseFloat(vcuRaw) : undefined,
rateLimit,
lastChecked: Date.now(), lastChecked: Date.now(),
}; };
} }
@ -119,15 +158,82 @@ export function onVeniceBalanceUpdate(
*/ */
export type BalanceStatus = "ok" | "low" | "critical" | "depleted" | "unknown"; 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( export function evaluateBalanceStatus(
balance: VeniceBalance | null, balance: VeniceBalance | null,
thresholds: VeniceBalanceThresholds = DEFAULT_VENICE_BALANCE_THRESHOLDS, thresholds: VeniceBalanceThresholds = DEFAULT_VENICE_BALANCE_THRESHOLDS,
): BalanceStatus { ): BalanceStatus {
if (!balance || balance.diem === undefined) return "unknown"; return evaluateBalanceDetailed(balance, thresholds).status;
if (balance.diem <= 0) return "depleted"; }
if (balance.diem < thresholds.criticalDiemThreshold) return "critical";
if (balance.diem < thresholds.lowDiemThreshold) return "low"; /**
return "ok"; * 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 (!thresholds.showInStatus) return null;
if (!balance) 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 diemLabel = balance.diem !== undefined ? balance.diem.toFixed(2) : "?";
const ageMs = Date.now() - balance.lastChecked; const ageMs = Date.now() - balance.lastChecked;
const ageLabel = formatAge(ageMs); const ageLabel = formatAge(ageMs);
@ -155,14 +261,26 @@ export function formatVeniceBalanceStatus(
const parts = [ const parts = [
`DIEM: ${diemLabel}`, `DIEM: ${diemLabel}`,
`Status: ${statusEmoji[status]} ${status.toUpperCase()}`,
`(${ageLabel})`,
]; ];
if (balance.usd !== undefined) { 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(" · "); return parts.join(" · ");
} }
@ -175,21 +293,39 @@ export function generateBalanceWarning(
thresholds: VeniceBalanceThresholds = DEFAULT_VENICE_BALANCE_THRESHOLDS, thresholds: VeniceBalanceThresholds = DEFAULT_VENICE_BALANCE_THRESHOLDS,
): string | null { ): string | null {
if (!thresholds.enabled) return null; if (!thresholds.enabled) return null;
if (!balance || balance.diem === undefined) return null; if (!balance) return null;
const status = evaluateBalanceStatus(balance, thresholds); const evaluation = evaluateBalanceDetailed(balance, thresholds);
const diemLabel = balance.diem.toFixed(2); const diemLabel = balance.diem !== undefined ? balance.diem.toFixed(2) : "?";
switch (status) { // DIEM-based warnings
case "critical": if (evaluation.reason === "diem") {
return `🚨 Venice balance critical: ${diemLabel} DIEM remaining. Consider topping up at https://venice.ai/settings/billing`; switch (evaluation.status) {
case "low": case "critical":
return `⚠️ Venice balance low: ${diemLabel} DIEM remaining`; return `🚨 Venice balance critical: ${diemLabel} DIEM remaining. Consider topping up at https://venice.ai/settings/billing`;
case "depleted": case "low":
return `❌ Venice balance depleted (0.00 DIEM). Top up at https://venice.ai/settings/billing`; return `⚠️ Venice balance low: ${diemLabel} DIEM remaining`;
default: case "depleted":
return null; 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, enabled: override.enabled ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.enabled,
lowDiemThreshold: override.lowDiemThreshold ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.lowDiemThreshold, lowDiemThreshold: override.lowDiemThreshold ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.lowDiemThreshold,
criticalDiemThreshold: override.criticalDiemThreshold ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.criticalDiemThreshold, 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, showInStatus: override.showInStatus ?? DEFAULT_VENICE_BALANCE_THRESHOLDS.showInStatus,
}; };
} }