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>

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(),
   })

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", () => {

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<string, string>):
   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,
   };
 }