Add /usage endpoint with token tracking, live balances, and cost calculations

.gitignore

@@ -1,3 +1,4 @@
 config.json
+usage-data.json
 node_modules/
 *.log

CLAUDE.md

@@ -5,8 +5,9 @@ Anthropic-compatible proxy that routes AI requests to multiple providers based o
 
 ## File Structure
 - `proxy.mjs` - Main proxy server (zero dependencies, Node.js built-ins only)
-- `config.json` - API keys and auth token (gitignored)
+- `config.json` - API keys, auth token, and pricing (gitignored)
 - `config.example.json` - Template config
+- `usage-data.json` - Token usage tracking, auto-generated and gitignored
 - `ein-ai-proxy.service` - Systemd unit file
 - `deploy.sh` - Deployment and testing script
 - `sync-context.sh` - Syncs global context to Claude Code and Roo Code
@@ -14,6 +15,7 @@ Anthropic-compatible proxy that routes AI requests to multiple providers based o
 ## Development
 Run locally: `node proxy.mjs` (requires `config.json`)
 Test health endpoint: `curl http://localhost:3100/health`
+Check usage: `curl http://localhost:3100/usage`
 Zero npm dependencies — do not add any.
 
 ## Deployment

CONTEXT.md

@@ -235,6 +235,9 @@ sudo journalctl -u ein-ai-proxy -f --no-pager
 # Check proxy health
 curl https://ai.ein-softworks.com/health
 
+# Check usage statistics
+curl https://ai.ein-softworks.com/usage
+
 # Pull and deploy changes
 cd ~/ai-proxy && git pull
 sudo systemctl restart ein-ai-proxy
@@ -268,6 +271,15 @@ GLM-4.7's thinking block signatures are different from Anthropic's and will be r
 ### Zero Dependencies
 The proxy has zero npm dependencies — it uses only Node.js built-ins (`node:http`, `node:fs`, etc.). Do not add dependencies unless absolutely necessary.
 
+### Usage Tracking
+The proxy automatically tracks token usage for all requests:
+- Token counts are extracted from both streaming and non-streaming responses
+- Data is persisted to `usage-data.json` (gitignored)
+- Monthly totals reset on the 1st of each month
+- Costs are calculated based on pricing in `config.json`
+- Live account balances are fetched from providers that support it (DeepSeek)
+- Use `curl https://ai.ein-softworks.com/usage` to view all usage data
+
 ## Key Reference URLs
 - **Proxy API**: https://ai.ein-softworks.com
 - **Gitea**: https://git.ein-softworks.com

README.md

@@ -59,7 +59,7 @@ cd ein-ai-proxy
 cp config.example.json config.json
 ```
 
-Edit `config.json` with your API keys:
+Edit `config.json` with your API keys and pricing:
 
 ```json
 {
@@ -69,10 +69,26 @@ Edit `config.json` with your API keys:
     "anthropic": { "apiKey": "sk-ant-..." },
     "zai":       { "apiKey": "your-zai-key" },
     "deepseek":  { "apiKey": "your-deepseek-key" }
+  },
+  "pricing": {
+    "anthropic": {
+      "inputPerMTok": 15,
+      "outputPerMTok": 75
+    },
+    "deepseek": {
+      "inputPerMTok": 0.27,
+      "outputPerMTok": 1.10
+    },
+    "zai": {
+      "type": "flat",
+      "monthlyCostUsd": 10
+    }
   }
 }
 ```
 
+**Pricing notes**: Rates are per million tokens (input/output). Update these values based on current provider pricing. Z.ai uses a flat-rate model ($10/mo) regardless of token usage.
+
 ### Run
 
 ```bash
@@ -185,6 +201,61 @@ curl https://ai.ein-softworks.com/health
 
 Returns the proxy status and current routing table. No authentication required.
 
+### Usage Statistics
+
+```bash
+curl https://ai.ein-softworks.com/usage
+```
+
+Returns comprehensive usage data grouped by provider:
+
+- **Per-model breakdown**: Requested model name, actual provider model, total + monthly token counts and request counts
+- **Cost calculations**: Dollar amounts for each model (based on pricing in `config.json`)
+- **Live balances**: Account balances from providers that support it (DeepSeek)
+- **Totals**: Aggregated request counts and costs across all models
+
+Example response:
+
+```json
+{
+  "service": "ein-ai-proxy",
+  "generatedAt": "2026-03-27T00:38:00.000Z",
+  "currentPeriod": "2026-03",
+  "providers": {
+    "anthropic": {
+      "balance": null,
+      "pricingNote": "No balance API available — check console.anthropic.com",
+      "models": {
+        "claude-opus-4-6": {
+          "routesTo": "claude-opus-4-6",
+          "total": { "inputTokens": 150000, "outputTokens": 50000, "requests": 42, "costUsd": 6.00 },
+          "monthly": { "inputTokens": 30000, "outputTokens": 10000, "requests": 8, "costUsd": 1.20 }
+        }
+      }
+    },
+    "deepseek": {
+      "balance": { "totalBalance": "8.52", "currency": "USD" },
+      "models": { ... }
+    },
+    "zai": {
+      "balance": null,
+      "pricing": "flat_rate",
+      "monthlyCostUsd": 10,
+      "models": { ... }
+    }
+  },
+  "totals": {
+    "totalRequests": 252,
+    "monthlyRequests": 53,
+    "totalCostUsd": 6.14,
+    "monthlyCostUsd": 1.24,
+    "note": "Cost excludes flat-rate models. Z.ai costs $10/mo regardless of usage."
+  }
+}
+```
+
+**Note**: Monthly totals reset on the 1st of each month. Usage data is persisted to `usage-data.json` in the proxy directory.
+
 ### Messages Endpoint
 
 ```bash

config.example.json

@@ -11,5 +11,19 @@
     "deepseek": {
       "apiKey": "YOUR_DEEPSEEK_KEY"
     }
+  },
+  "pricing": {
+    "anthropic": {
+      "inputPerMTok": 15,
+      "outputPerMTok": 75
+    },
+    "deepseek": {
+      "inputPerMTok": 0.27,
+      "outputPerMTok": 1.10
+    },
+    "zai": {
+      "type": "flat",
+      "monthlyCostUsd": 10
+    }
   }
 }

proxy.mjs

@@ -1,5 +1,5 @@
 import { createServer } from "node:http";
-import { readFileSync, existsSync } from "node:fs";
+import { readFileSync, existsSync, writeFileSync, renameSync, mkdirSync } from "node:fs";
 import { join, dirname } from "node:path";
 import { fileURLToPath } from "node:url";
 
@@ -8,6 +8,7 @@ const __dirname = dirname(fileURLToPath(import.meta.url));
 // --- Configuration -----------------------------------------------------------
 
 const CONFIG_PATH = join(__dirname, "config.json");
+const USAGE_DATA_PATH = join(__dirname, "usage-data.json");
 
 function loadConfig() {
   if (!existsSync(CONFIG_PATH)) {
@@ -22,6 +23,13 @@ const config = loadConfig();
 const PORT = config.port || 3100;
 const AUTH_TOKEN = config.authToken;
 
+// Default pricing if not in config
+const DEFAULT_PRICING = {
+  anthropic: { inputPerMTok: 15, outputPerMTok: 75 }, // Opus rates
+  deepseek: { inputPerMTok: 0.27, outputPerMTok: 1.10 },
+  zai: { type: "flat", monthlyCostUsd: 10 },
+};
+
 // Provider definitions: maps provider name -> { baseUrl, apiKey, authStyle }
 // authStyle: "anthropic" sends x-api-key header, "bearer" sends Authorization: Bearer
 const PROVIDERS = {
@@ -69,6 +77,248 @@ const MODEL_ROUTES = {
 // Default route for unknown models
 const DEFAULT_ROUTE = { provider: "zai", rewriteModel: "glm-4.7" };
 
+// --- Usage Tracking -----------------------------------------------------------
+
+function getCurrentPeriod() {
+  const now = new Date();
+  return `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, "0")}`;
+}
+
+function loadUsageData() {
+  if (!existsSync(USAGE_DATA_PATH)) {
+    return {
+      version: 1,
+      lastUpdated: new Date().toISOString(),
+      models: {},
+    };
+  }
+  try {
+    return JSON.parse(readFileSync(USAGE_DATA_PATH, "utf-8"));
+  } catch (err) {
+    console.error(`[usage] Failed to load usage data: ${err.message}`);
+    return {
+      version: 1,
+      lastUpdated: new Date().toISOString(),
+      models: {},
+    };
+  }
+}
+
+function saveUsageData(data) {
+  try {
+    const tempPath = USAGE_DATA_PATH + ".tmp";
+    writeFileSync(tempPath, JSON.stringify(data, null, 2), "utf-8");
+    renameSync(tempPath, USAGE_DATA_PATH);
+  } catch (err) {
+    console.error(`[usage] Failed to save usage data: ${err.message}`);
+  }
+}
+
+function recordUsage(requestedModel, provider, actualModel, inputTokens, outputTokens) {
+  const usageData = loadUsageData();
+  const currentPeriod = getCurrentPeriod();
+
+  if (!usageData.models[requestedModel]) {
+    usageData.models[requestedModel] = {
+      provider,
+      actualModel,
+      total: { inputTokens: 0, outputTokens: 0, requests: 0 },
+      monthly: { period: currentPeriod, inputTokens: 0, outputTokens: 0, requests: 0 },
+    };
+  }
+
+  const modelData = usageData.models[requestedModel];
+
+  // Check if monthly period has rolled over
+  if (modelData.monthly.period !== currentPeriod) {
+    modelData.monthly = { period: currentPeriod, inputTokens: 0, outputTokens: 0, requests: 0 };
+  }
+
+  // Update totals
+  modelData.total.inputTokens += inputTokens || 0;
+  modelData.total.outputTokens += outputTokens || 0;
+  modelData.total.requests += 1;
+
+  // Update monthly
+  modelData.monthly.inputTokens += inputTokens || 0;
+  modelData.monthly.outputTokens += outputTokens || 0;
+  modelData.monthly.requests += 1;
+
+  usageData.lastUpdated = new Date().toISOString();
+  saveUsageData(usageData);
+}
+
+// --- Provider Balance Fetching ----------------------------------------------
+
+async function fetchDeepSeekBalance(apiKey) {
+  try {
+    const res = await fetch("https://api.deepseek.com/user/balance", {
+      headers: { "Authorization": `Bearer ${apiKey}` },
+      signal: AbortSignal.timeout(5000),
+    });
+    if (res.ok) {
+      const data = await res.json();
+      if (data.balance_infos && data.balance_infos.length > 0) {
+        const balance = data.balance_infos.find((b) => b.currency === "USD") || data.balance_infos[0];
+        return {
+          totalBalance: balance.total_balance || balance.total_granted_balance,
+          currency: balance.currency || "USD",
+        };
+      }
+    }
+    return null;
+  } catch (err) {
+    console.error(`[usage] Failed to fetch DeepSeek balance: ${err.message}`);
+    return null;
+  }
+}
+
+async function fetchAnthropicBalance(apiKey) {
+  // Anthropic doesn't have a public balance API for regular API keys
+  // Admin API exists but requires a different key type and org access
+  return null;
+}
+
+async function fetchZaiBalance(apiKey) {
+  // Z.ai doesn't have a documented balance API; flat-rate subscription
+  return null;
+}
+
+async function fetchProviderBalances() {
+  const balances = {};
+  const promises = [];
+
+  if (PROVIDERS.deepseek?.apiKey) {
+    promises.push(
+      fetchDeepSeekBalance(PROVIDERS.deepseek.apiKey).then((b) => {
+        if (b) balances.deepseek = b;
+      })
+    );
+  }
+  if (PROVIDERS.anthropic?.apiKey) {
+    promises.push(
+      fetchAnthropicBalance(PROVIDERS.anthropic.apiKey).then((b) => {
+        if (b) balances.anthropic = b;
+      })
+    );
+  }
+  if (PROVIDERS.zai?.apiKey) {
+    promises.push(
+      fetchZaiBalance(PROVIDERS.zai.apiKey).then((b) => {
+        if (b) balances.zai = b;
+      })
+    );
+  }
+
+  await Promise.allSettled(promises);
+  return balances;
+}
+
+// --- Cost Calculation --------------------------------------------------------
+
+function calculateCost(provider, inputTokens, outputTokens) {
+  const pricing = config.pricing?.[provider] || DEFAULT_PRICING[provider];
+  if (!pricing || pricing.type === "flat") {
+    return null;
+  }
+  const inputCost = ((inputTokens || 0) / 1_000_000) * pricing.inputPerMTok;
+  const outputCost = ((outputTokens || 0) / 1_000_000) * pricing.outputPerMTok;
+  return inputCost + outputCost;
+}
+
+// --- Usage Data Aggregation -------------------------------------------------
+
+function buildUsageResponse() {
+  const usageData = loadUsageData();
+  const balances = fetchProviderBalances(); // Start fetching but don't await
+  const currentPeriod = getCurrentPeriod();
+
+  const providers = {
+    anthropic: { balance: null, models: {} },
+    deepseek: { balance: null, models: {} },
+    zai: { balance: null, models: {} },
+  };
+
+  let totalRequests = 0;
+  let monthlyRequests = 0;
+  let totalCost = 0;
+  let monthlyCost = 0;
+  let flatRateModels = 0;
+
+  // Process model data
+  for (const [requestedModel, modelData] of Object.entries(usageData.models || {})) {
+    const providerName = modelData.provider;
+    if (!providers[providerName]) continue;
+
+    const { total, monthly, actualModel } = modelData;
+
+    providers[providerName].models[requestedModel] = {
+      routesTo: actualModel || requestedModel,
+      total: {
+        inputTokens: total.inputTokens,
+        outputTokens: total.outputTokens,
+        requests: total.requests,
+        costUsd: calculateCost(providerName, total.inputTokens, total.outputTokens),
+      },
+      monthly: {
+        period: monthly.period,
+        inputTokens: monthly.inputTokens,
+        outputTokens: monthly.outputTokens,
+        requests: monthly.requests,
+        costUsd: calculateCost(providerName, monthly.inputTokens, monthly.outputTokens),
+      },
+    };
+
+    totalRequests += total.requests;
+    monthlyRequests += monthly.requests;
+
+    const totalCostVal = calculateCost(providerName, total.inputTokens, total.outputTokens);
+    const monthlyCostVal = calculateCost(providerName, monthly.inputTokens, monthly.outputTokens);
+
+    if (totalCostVal !== null) totalCost += totalCostVal;
+    else flatRateModels++;
+
+    if (monthlyCostVal !== null) monthlyCost += monthlyCostVal;
+  }
+
+  // Calculate pricing per provider
+  for (const [providerName, providerData] of Object.entries(providers)) {
+    const pricing = config.pricing?.[providerName] || DEFAULT_PRICING[providerName];
+    if (pricing?.type === "flat") {
+      providerData.pricing = "flat_rate";
+      providerData.monthlyCostUsd = pricing.monthlyCostUsd;
+    } else {
+      providerData.pricing = {
+        inputPerMTok: pricing?.inputPerMTok,
+        outputPerMTok: pricing?.outputPerMTok,
+      };
+    }
+  }
+
+  // Build note
+  const notes = [];
+  if (flatRateModels > 0) {
+    notes.push(`Cost excludes ${flatRateModels} flat-rate model(s). Z.ai costs $10/mo regardless of usage.`);
+  }
+  if (!balances.anthropic) {
+    notes.push("Anthropic has no public balance API — check console.anthropic.com");
+  }
+  if (!balances.zai) {
+    notes.push("Z.ai is flat-rate; no balance API available");
+  }
+
+  return {
+    usageData,
+    providers,
+    totalRequests,
+    monthlyRequests,
+    totalCost,
+    monthlyCost,
+    currentPeriod,
+    notes,
+  };
+}
+
 // --- Helpers -----------------------------------------------------------------
 
 function resolveRoute(model) {
@@ -109,6 +359,44 @@ function buildUpstreamHeaders(provider, originalHeaders) {
 
 // --- Request Handler ---------------------------------------------------------
 
+// Extract usage from SSE stream data
+function extractUsageFromSSE(text, provider) {
+  let inputTokens = 0;
+  let outputTokens = 0;
+
+  // Split by "data:" lines (SSE format)
+  const lines = text.split("\n");
+  for (const line of lines) {
+    const trimmed = line.trim();
+    if (trimmed.startsWith("event: message_start")) {
+      // Find the next data line
+      const dataIndex = lines.indexOf(line) + 1;
+      if (dataIndex < lines.length && lines[dataIndex].trim().startsWith("data:")) {
+        try {
+          const dataStr = lines[dataIndex].trim().slice(5).trim();
+          const data = JSON.parse(dataStr);
+          if (data.message?.usage?.input_tokens) {
+            inputTokens = data.message.usage.input_tokens;
+          }
+        } catch {}
+      }
+    } else if (trimmed.startsWith("event: message_delta")) {
+      const dataIndex = lines.indexOf(line) + 1;
+      if (dataIndex < lines.length && lines[dataIndex].trim().startsWith("data:")) {
+        try {
+          const dataStr = lines[dataIndex].trim().slice(5).trim();
+          const data = JSON.parse(dataStr);
+          if (data.usage?.output_tokens) {
+            outputTokens = data.usage.output_tokens;
+          }
+        } catch {}
+      }
+    }
+  }
+
+  return { inputTokens, outputTokens };
+}
+
 async function handleMessages(req, res) {
   // Read body
   const chunks = [];
@@ -135,6 +423,7 @@ async function handleMessages(req, res) {
   }
 
   // Rewrite model name if needed
+  const actualModel = route.rewriteModel || body.model;
   if (route.rewriteModel) {
     body.model = route.rewriteModel;
   }
@@ -158,7 +447,7 @@ async function handleMessages(req, res) {
     const statusCode = upstreamRes.status;
 
     if (isStreaming && statusCode === 200) {
-      // SSE streaming: pipe through directly
+      // SSE streaming: pipe through directly while tracking usage
       res.writeHead(statusCode, {
         "content-type": "text/event-stream",
         "cache-control": "no-cache",
@@ -168,14 +457,22 @@ async function handleMessages(req, res) {
 
       const reader = upstreamRes.body.getReader();
       const decoder = new TextDecoder();
+      let fullText = "";
 
       try {
         while (true) {
           const { done, value } = await reader.read();
           if (done) break;
           const text = decoder.decode(value, { stream: true });
+          fullText += text;
           res.write(text);
         }
+
+        // Extract usage from the accumulated SSE text
+        const { inputTokens, outputTokens } = extractUsageFromSSE(fullText, route.provider);
+        if (inputTokens > 0 || outputTokens > 0) {
+          recordUsage(requestedModel, route.provider, actualModel, inputTokens, outputTokens);
+        }
       } catch (streamErr) {
         console.error(`[proxy] Stream error from ${route.provider}:`, streamErr.message);
       } finally {
@@ -185,6 +482,23 @@ async function handleMessages(req, res) {
       // Non-streaming or error: forward the full response
       const responseBody = await upstreamRes.text();
       const responseHeaders = { "content-type": upstreamRes.headers.get("content-type") || "application/json" };
+
+      // Try to extract usage from non-streaming response
+      if (statusCode === 200 && responseHeaders["content-type"]?.includes("application/json")) {
+        try {
+          const json = JSON.parse(responseBody);
+          if (json.usage?.input_tokens || json.usage?.output_tokens) {
+            recordUsage(
+              requestedModel,
+              route.provider,
+              actualModel,
+              json.usage.input_tokens,
+              json.usage.output_tokens
+            );
+          }
+        } catch {}
+      }
+
       res.writeHead(statusCode, responseHeaders);
       res.end(responseBody);
     }
@@ -234,6 +548,50 @@ const server = createServer(async (req, res) => {
     return;
   }
 
+  // Usage endpoint (no auth required - no sensitive data)
+  if (req.method === "GET" && urlPath === "/usage") {
+    const { usageData, providers, totalRequests, monthlyRequests, totalCost, monthlyCost, currentPeriod, notes } =
+      buildUsageResponse();
+
+    // Fetch live balances in parallel
+    const balances = await fetchProviderBalances();
+
+    // Apply balances to provider data
+    for (const [providerName, balance] of Object.entries(balances)) {
+      if (providers[providerName]) {
+        providers[providerName].balance = balance;
+      }
+    }
+
+    // Add balance notes
+    if (!balances.anthropic) {
+      providers.anthropic.balance = null;
+      providers.anthropic.pricingNote = "No balance API available — check console.anthropic.com";
+    }
+    if (!balances.zai) {
+      providers.zai.balance = null;
+      providers.zai.pricingNote = "Flat-rate subscription, no balance API available";
+    }
+
+    res.writeHead(200, { "content-type": "application/json" });
+    res.end(
+      JSON.stringify({
+        service: "ein-ai-proxy",
+        generatedAt: new Date().toISOString(),
+        currentPeriod,
+        providers,
+        totals: {
+          totalRequests,
+          monthlyRequests,
+          totalCostUsd: Number(totalCost.toFixed(4)),
+          monthlyCostUsd: Number(monthlyCost.toFixed(4)),
+          note: notes.join(" ") || null,
+        },
+      })
+    );
+    return;
+  }
+
   // Auth check (everything below requires auth)
   if (AUTH_TOKEN) {
     const authHeader = req.headers["authorization"] || "";