Files
ein-ai-proxy/proxy.mjs
T
2026-03-27 11:04:42 -05:00

735 lines
25 KiB
JavaScript

import { createServer } from "node:http";
import { readFileSync, existsSync, writeFileSync, renameSync, mkdirSync } from "node:fs";
import { join, dirname } from "node:path";
import { fileURLToPath } from "node:url";
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)) {
console.error(`Config file not found: ${CONFIG_PATH}`);
console.error("Copy config.example.json to config.json and fill in your keys.");
process.exit(1);
}
return JSON.parse(readFileSync(CONFIG_PATH, "utf-8"));
}
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 = {
anthropic: {
baseUrl: "https://api.anthropic.com",
apiKey: config.providers.anthropic.apiKey,
authStyle: "api-key",
},
zai: {
baseUrl: "https://api.z.ai/api/anthropic",
apiKey: config.providers.zai.apiKey,
authStyle: "bearer",
},
deepseek: {
baseUrl: "https://api.deepseek.com/anthropic",
apiKey: config.providers.deepseek.apiKey,
authStyle: "bearer",
},
};
// Model -> provider routing table
// Claude Code sends model names like "claude-sonnet-4-6", "claude-opus-4-6", "claude-haiku-4-5-20251001"
// We map these to the correct provider + optionally rewrite the model name
const MODEL_ROUTES = {
// Opus models -> Anthropic (the expensive brain)
"claude-opus-4-6": { provider: "anthropic" },
"claude-opus-4-5": { provider: "anthropic" },
"claude-opus-4-5-20251101": { provider: "anthropic" },
// Sonnet models -> Z.ai GLM-4.7 (the workhorse)
"claude-sonnet-4-6": { provider: "zai", rewriteModel: "glm-4.7" },
"claude-sonnet-4-5": { provider: "zai", rewriteModel: "glm-4.7" },
"claude-sonnet-4-5-20250929": { provider: "zai", rewriteModel: "glm-4.7" },
// Haiku models -> DeepSeek (the cheap tier)
"claude-haiku-4-5-20251001": { provider: "deepseek", rewriteModel: "deepseek-chat" },
"claude-haiku-4-5": { provider: "deepseek", rewriteModel: "deepseek-chat" },
// Direct model names (when client explicitly asks for a provider's model)
"glm-4.7": { provider: "zai" },
"glm-4.5-air": { provider: "zai" },
"deepseek-chat": { provider: "deepseek" },
};
// Default route for unknown models
const DEFAULT_ROUTE = { provider: "zai", rewriteModel: "glm-4.7" };
// --- Provider Status (in-memory, resets on restart) --------------------------
const providerStatus = {};
for (const name of Object.keys(PROVIDERS)) {
providerStatus[name] = { status: "idle", lastSeen: null, lastError: null };
}
// --- Usage Tracking -----------------------------------------------------------
function getCurrentPeriod() {
const now = new Date();
return `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, "0")}`;
}
function getCurrentDate() {
return new Date().toISOString().slice(0, 10); // "2026-03-27"
}
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();
const todayStr = getCurrentDate();
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 },
daily: { date: todayStr, 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 };
}
// Check if daily date has rolled over (or missing for backward compat)
if (!modelData.daily || modelData.daily.date !== todayStr) {
modelData.daily = { date: todayStr, 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;
// Update daily
modelData.daily.inputTokens += inputTokens || 0;
modelData.daily.outputTokens += outputTokens || 0;
modelData.daily.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 getProviderPricing(providerName) {
return config.pricing?.[providerName] || DEFAULT_PRICING[providerName];
}
function buildModelPricing(providerName) {
const pricing = getProviderPricing(providerName);
if (pricing?.type === "flat") return "flat_rate";
return { inputPerMTok: pricing?.inputPerMTok, outputPerMTok: pricing?.outputPerMTok };
}
function buildUsageResponse() {
const usageData = loadUsageData();
const currentPeriod = getCurrentPeriod();
const todayStr = getCurrentDate();
const providers = {
anthropic: { balance: null, models: {} },
deepseek: { balance: null, models: {} },
zai: { balance: null, models: {} },
};
let totalRequests = 0;
let monthlyRequests = 0;
let todayRequests = 0;
let totalCost = 0;
let monthlyCost = 0;
let todayCost = 0;
let totalInputTokens = 0;
let totalOutputTokens = 0;
let monthlyInputTokens = 0;
let monthlyOutputTokens = 0;
let todayInputTokens = 0;
let todayOutputTokens = 0;
let flatRateProviders = new Set();
// Process model data from usage records
for (const [requestedModel, modelData] of Object.entries(usageData.models || {})) {
const providerName = modelData.provider;
if (!providers[providerName]) continue;
const { total, monthly, daily, actualModel } = modelData;
// Daily data: only include if it's from today, otherwise zeros
const todayData = (daily && daily.date === todayStr) ? daily : { inputTokens: 0, outputTokens: 0, requests: 0 };
const todayCostVal = calculateCost(providerName, todayData.inputTokens, todayData.outputTokens);
providers[providerName].models[requestedModel] = {
routesTo: actualModel || requestedModel,
pricing: buildModelPricing(providerName),
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),
},
today: {
date: todayStr,
inputTokens: todayData.inputTokens,
outputTokens: todayData.outputTokens,
requests: todayData.requests,
costUsd: todayCostVal,
},
};
// Aggregate totals
totalRequests += total.requests;
totalInputTokens += total.inputTokens;
totalOutputTokens += total.outputTokens;
monthlyRequests += monthly.period === currentPeriod ? monthly.requests : 0;
monthlyInputTokens += monthly.period === currentPeriod ? monthly.inputTokens : 0;
monthlyOutputTokens += monthly.period === currentPeriod ? monthly.outputTokens : 0;
todayRequests += todayData.requests;
todayInputTokens += todayData.inputTokens;
todayOutputTokens += todayData.outputTokens;
const totalCostVal = calculateCost(providerName, total.inputTokens, total.outputTokens);
const monthlyCostVal = calculateCost(providerName, monthly.period === currentPeriod ? monthly.inputTokens : 0, monthly.period === currentPeriod ? monthly.outputTokens : 0);
if (totalCostVal !== null) totalCost += totalCostVal;
else flatRateProviders.add(providerName);
if (monthlyCostVal !== null) monthlyCost += monthlyCostVal;
if (todayCostVal !== null) todayCost += todayCostVal;
}
// Seed all configured models from MODEL_ROUTES (ensures zero-usage models appear)
for (const [requestedModel, route] of Object.entries(MODEL_ROUTES)) {
const providerName = route.provider;
if (!providers[providerName]) continue;
// Skip if already populated from usage data
if (providers[providerName].models[requestedModel]) continue;
providers[providerName].models[requestedModel] = {
routesTo: route.rewriteModel || requestedModel,
pricing: buildModelPricing(providerName),
total: { inputTokens: 0, outputTokens: 0, requests: 0, costUsd: calculateCost(providerName, 0, 0) },
monthly: { period: currentPeriod, inputTokens: 0, outputTokens: 0, requests: 0, costUsd: calculateCost(providerName, 0, 0) },
today: { date: todayStr, inputTokens: 0, outputTokens: 0, requests: 0, costUsd: calculateCost(providerName, 0, 0) },
};
}
// Set provider-level metadata: pricing, status
for (const [providerName, providerData] of Object.entries(providers)) {
const pricing = getProviderPricing(providerName);
if (pricing?.type === "flat") {
providerData.pricing = "flat_rate";
providerData.monthlyCostUsd = pricing.monthlyCostUsd;
} else {
providerData.pricing = {
inputPerMTok: pricing?.inputPerMTok,
outputPerMTok: pricing?.outputPerMTok,
};
}
// Provider status from in-memory tracking
const status = providerStatus[providerName] || { status: "idle", lastSeen: null, lastError: null };
providerData.status = status.status;
providerData.lastSeen = status.lastSeen;
providerData.lastError = status.lastError;
}
// Build notes (balance notes are added in the /usage handler after balances resolve)
const notes = [];
if (flatRateProviders.size > 0) {
notes.push("Cost excludes flat-rate provider(s). Z.ai costs $10/mo regardless of usage.");
}
return {
providers,
currentPeriod,
notes,
totals: {
totalRequests,
monthlyRequests,
todayRequests,
totalInputTokens,
totalOutputTokens,
monthlyInputTokens,
monthlyOutputTokens,
todayInputTokens,
todayOutputTokens,
totalCostUsd: Number(totalCost.toFixed(4)),
monthlyCostUsd: Number(monthlyCost.toFixed(4)),
todayCostUsd: Number(todayCost.toFixed(4)),
},
};
}
// --- Helpers -----------------------------------------------------------------
function resolveRoute(model) {
// Exact match first
if (MODEL_ROUTES[model]) return MODEL_ROUTES[model];
// Partial match: check if model starts with a known prefix
for (const [pattern, route] of Object.entries(MODEL_ROUTES)) {
if (model.startsWith(pattern)) return route;
}
// Fallback
console.log(`[route] Unknown model "${model}", using default -> zai/glm-4.7`);
return DEFAULT_ROUTE;
}
function buildUpstreamHeaders(provider, originalHeaders) {
const headers = {
"content-type": "application/json",
"anthropic-version": originalHeaders["anthropic-version"] || "2023-06-01",
};
// Copy through any anthropic-beta headers
if (originalHeaders["anthropic-beta"]) {
headers["anthropic-beta"] = originalHeaders["anthropic-beta"];
}
// Auth style depends on provider
const prov = PROVIDERS[provider];
if (prov.authStyle === "api-key") {
headers["x-api-key"] = prov.apiKey;
} else {
headers["authorization"] = `Bearer ${prov.apiKey}`;
}
return headers;
}
// --- 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 = [];
for await (const chunk of req) chunks.push(chunk);
const rawBody = Buffer.concat(chunks).toString("utf-8");
let body;
try {
body = JSON.parse(rawBody);
} catch {
res.writeHead(400, { "content-type": "application/json" });
res.end(JSON.stringify({ error: { type: "invalid_request", message: "Invalid JSON body" } }));
return;
}
const requestedModel = body.model || "claude-sonnet-4-6";
const route = resolveRoute(requestedModel);
const provider = PROVIDERS[route.provider];
if (!provider) {
res.writeHead(500, { "content-type": "application/json" });
res.end(JSON.stringify({ error: { type: "config_error", message: `Unknown provider: ${route.provider}` } }));
return;
}
// Rewrite model name if needed
const actualModel = route.rewriteModel || body.model;
if (route.rewriteModel) {
body.model = route.rewriteModel;
}
const isStreaming = body.stream === true;
const upstreamUrl = `${provider.baseUrl}/v1/messages`;
const upstreamHeaders = buildUpstreamHeaders(route.provider, req.headers);
console.log(
`[proxy] ${requestedModel} -> ${route.provider}/${body.model} (stream=${isStreaming})`
);
try {
const upstreamRes = await fetch(upstreamUrl, {
method: "POST",
headers: upstreamHeaders,
body: JSON.stringify(body),
});
// Forward status code
const statusCode = upstreamRes.status;
if (isStreaming && statusCode === 200) {
// SSE streaming: pipe through directly while tracking usage
res.writeHead(statusCode, {
"content-type": "text/event-stream",
"cache-control": "no-cache",
"connection": "keep-alive",
"access-control-allow-origin": "*",
});
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);
}
providerStatus[route.provider] = { status: "ok", lastSeen: new Date().toISOString(), lastError: null };
} catch (streamErr) {
console.error(`[proxy] Stream error from ${route.provider}:`, streamErr.message);
providerStatus[route.provider] = { status: "error", lastSeen: new Date().toISOString(), lastError: streamErr.message };
} finally {
res.end();
}
} else {
// 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 {}
}
// Update provider status based on HTTP status
if (statusCode >= 200 && statusCode < 300) {
providerStatus[route.provider] = { status: "ok", lastSeen: new Date().toISOString(), lastError: null };
} else {
providerStatus[route.provider] = { status: "degraded", lastSeen: new Date().toISOString(), lastError: `HTTP ${statusCode}` };
}
res.writeHead(statusCode, responseHeaders);
res.end(responseBody);
}
} catch (fetchErr) {
console.error(`[proxy] Fetch error to ${route.provider} (${upstreamUrl}):`, fetchErr.message);
providerStatus[route.provider] = { status: "error", lastSeen: new Date().toISOString(), lastError: fetchErr.message };
res.writeHead(502, { "content-type": "application/json" });
res.end(
JSON.stringify({
error: {
type: "proxy_error",
message: `Failed to reach upstream provider ${route.provider}: ${fetchErr.message}`,
},
})
);
}
}
// --- Server ------------------------------------------------------------------
const server = createServer(async (req, res) => {
// CORS preflight
if (req.method === "OPTIONS") {
res.writeHead(204, {
"access-control-allow-origin": "*",
"access-control-allow-methods": "POST, OPTIONS",
"access-control-allow-headers": "content-type, authorization, x-api-key, anthropic-version, anthropic-beta",
});
res.end();
return;
}
const urlPath = req.url.split("?")[0];
// Health check (no auth required)
if (req.method === "GET" && (urlPath === "/" || urlPath === "/health")) {
res.writeHead(200, { "content-type": "application/json" });
res.end(
JSON.stringify({
status: "ok",
service: "ein-ai-proxy",
providers: Object.keys(PROVIDERS),
routes: Object.fromEntries(
Object.entries(MODEL_ROUTES).map(([m, r]) => [m, `${r.provider}/${r.rewriteModel || m}`])
),
})
);
return;
}
// Usage endpoint (no auth required - no sensitive data)
if (req.method === "GET" && urlPath === "/usage") {
const { providers, currentPeriod, notes, totals } = 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";
}
notes.push("Anthropic has no public balance API — check console.anthropic.com");
notes.push("Z.ai is flat-rate; no balance API available");
// Build routing summary (the 3 Claude Code / Roo Code tiers)
const routing = {
sonnet: { provider: "zai", model: "glm-4.7", description: "Daily workhorse (Z.ai flat-rate)" },
opus: { provider: "anthropic", model: "claude-opus-4-6", description: "Complex tasks (Anthropic pay-per-token)" },
haiku: { provider: "deepseek", model: "deepseek-chat", description: "Trivial/fast tasks (DeepSeek cheap tier)" },
default: { provider: "zai", model: "glm-4.7", description: "Unknown models fall back here" },
};
res.writeHead(200, { "content-type": "application/json" });
res.end(
JSON.stringify({
service: "ein-ai-proxy",
generatedAt: new Date().toISOString(),
currentPeriod,
routing,
providers,
totals: {
...totals,
note: notes.join(" ") || null,
},
})
);
return;
}
// Auth check (everything below requires auth)
if (AUTH_TOKEN) {
const authHeader = req.headers["authorization"] || "";
const apiKeyHeader = req.headers["x-api-key"] || "";
const token = authHeader.replace(/^Bearer\s+/i, "") || apiKeyHeader;
if (token !== AUTH_TOKEN) {
res.writeHead(401, { "content-type": "application/json" });
res.end(JSON.stringify({ error: { type: "authentication_error", message: "Invalid auth token" } }));
return;
}
}
// Anthropic Messages API
if (req.method === "POST" && urlPath === "/v1/messages") {
await handleMessages(req, res);
return;
}
// 404 for everything else
res.writeHead(404, { "content-type": "application/json" });
res.end(JSON.stringify({ error: { type: "not_found", message: `Unknown path: ${urlPath}` } }));
});
server.listen(PORT, "0.0.0.0", () => {
console.log(`[ein-ai-proxy] Listening on http://0.0.0.0:${PORT}`);
console.log(`[ein-ai-proxy] Routes:`);
for (const [model, route] of Object.entries(MODEL_ROUTES)) {
console.log(` ${model} -> ${route.provider}/${route.rewriteModel || model}`);
}
console.log(` (default) -> zai/glm-4.7`);
});