ein-ai-proxy/proxy.mjs

import { createServer } from "node:http";
import { readFileSync, existsSync } 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");

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;

// 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" };

// --- 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 ---------------------------------------------------------

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
  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
      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();

      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
          const text = decoder.decode(value, { stream: true });
          res.write(text);
        }
      } catch (streamErr) {
        console.error(`[proxy] Stream error from ${route.provider}:`, 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" };
      res.writeHead(statusCode, responseHeaders);
      res.end(responseBody);
    }
  } catch (fetchErr) {
    console.error(`[proxy] Fetch error to ${route.provider} (${upstreamUrl}):`, 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;
  }

  // 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`);
});