From e7fe5732b703f5fb4c5a609655f3b3bf0d572963 Mon Sep 17 00:00:00 2001 From: William Stuckey Date: Thu, 26 Mar 2026 19:43:34 -0500 Subject: [PATCH] Add /usage endpoint with token tracking, live balances, and cost calculations --- .gitignore | 1 + CLAUDE.md | 4 +- CONTEXT.md | 12 ++ README.md | 73 ++++++++- config.example.json | 14 ++ proxy.mjs | 362 +++++++++++++++++++++++++++++++++++++++++++- 6 files changed, 462 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index dc1d633..7b0e54a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ config.json +usage-data.json node_modules/ *.log diff --git a/CLAUDE.md b/CLAUDE.md index 2729931..20903c6 100644 --- a/CLAUDE.md +++ b/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 diff --git a/CONTEXT.md b/CONTEXT.md index 5f5655a..3f4789d 100644 --- a/CONTEXT.md +++ b/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 diff --git a/README.md b/README.md index 976db3b..076295e 100644 --- a/README.md +++ b/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 diff --git a/config.example.json b/config.example.json index ea79d3a..dc712e5 100644 --- a/config.example.json +++ b/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 + } } } diff --git a/proxy.mjs b/proxy.mjs index b9be5d5..968583a 100644 --- a/proxy.mjs +++ b/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"] || "";