diff --git a/.gitignore b/.gitignore index 852f2a4..02fb6b0 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ node_modules/ .env .DS_Store *.txt +server-key.json diff --git a/key-manager.js b/key-manager.js index ae552ab..9031527 100644 --- a/key-manager.js +++ b/key-manager.js @@ -10,7 +10,9 @@ class KeyManager { key: key, success: 0, fail: 0, - deprecated: false // 是否被废弃 + deprecated: false, // 是否被废弃 + depleted: false, // 是否余额用尽 + balance: null // 最近一次余额信息 { totalAllowance, used, remaining, fetchedAt } })); this.algorithm = algorithm; // 'weighted' or 'simple' @@ -18,6 +20,7 @@ class KeyManager { this.endpointStats = {}; // 端点统计 { endpoint: { success: 0, fail: 0 } } this.removeOn402 = removeOn402; // 是否在402时移除key this.deprecatedKeys = []; // 已废弃的key列表 + this.skipThreshold = 0; // 剩余≤阈值则视为用尽 logInfo(`KeyManager initialized with ${this.keys.length} keys, algorithm: ${this.algorithm}, removeOn402: ${this.removeOn402}`); } @@ -27,7 +30,7 @@ class KeyManager { */ selectKey() { // 获取未废弃的key - const activeKeys = this.keys.filter(k => !k.deprecated); + const activeKeys = this.keys.filter(k => !k.deprecated && !k.depleted); if (activeKeys.length === 0) { throw new Error('No active keys available - all keys have been deprecated'); @@ -126,6 +129,38 @@ class KeyManager { logDebug(`Endpoint stats updated: ${endpoint}, success=${this.endpointStats[endpoint].success}, fail=${this.endpointStats[endpoint].fail}`); } + /** + * 更新余额与用尽状态（按索引） + * @param {number} index + * @param {object} balanceInfo { totalAllowance, used, remaining, fetchedAt } + */ + setBalanceByIndex(index, balanceInfo) { + if (typeof index !== 'number' || index < 0 || index >= this.keys.length) return; + const k = this.keys[index]; + if (!k) return; + k.balance = balanceInfo || null; + if (balanceInfo && typeof balanceInfo.remaining === 'number') { + k.depleted = balanceInfo.remaining <= this.skipThreshold; + } + } + + /** + * 设置跳过阈值（剩余≤n 视为用尽）并重算状态 + */ + setSkipThreshold(n) { + const val = Number(n); + this.skipThreshold = Number.isFinite(val) && val >= 0 ? val : 0; + for (const k of this.keys) { + if (k.balance && typeof k.balance.remaining === 'number') { + k.depleted = k.balance.remaining <= this.skipThreshold; + } + } + } + + getSkipThreshold() { + return this.skipThreshold; + } + /** * 废弃一个key * @param {string} key - 要废弃的key @@ -167,9 +202,10 @@ class KeyManager { * 获取统计信息 */ getStats() { - // 分离活跃的key和废弃的key - const activeKeys = this.keys.filter(k => !k.deprecated); - const deprecatedKeys = this.keys.filter(k => k.deprecated); + // 分离活跃的key和废弃的key（保留原始索引以便在UI中定位） + const indexedKeys = this.keys.map((k, idx) => ({ ...k, __index: idx })); + const activeKeys = indexedKeys.filter(k => !k.deprecated); + const deprecatedKeys = indexedKeys.filter(k => k.deprecated); return { algorithm: this.algorithm, @@ -181,7 +217,9 @@ class KeyManager { total: keyObj.success + keyObj.fail, successRate: keyObj.success + keyObj.fail > 0 ? ((keyObj.success / (keyObj.success + keyObj.fail)) * 100).toFixed(2) + '%' - : 'N/A' + : 'N/A', + index: keyObj.__index, + depleted: !!keyObj.depleted })), deprecatedKeys: deprecatedKeys.map(keyObj => ({ key: this.maskKey(keyObj.key), @@ -191,7 +229,8 @@ class KeyManager { successRate: keyObj.success + keyObj.fail > 0 ? ((keyObj.success / (keyObj.success + keyObj.fail)) * 100).toFixed(2) + '%' : 'N/A', - deprecatedAt: this.deprecatedKeys.find(dk => dk.key === keyObj.key)?.deprecatedAt || 'Unknown' + deprecatedAt: this.deprecatedKeys.find(dk => dk.key === keyObj.key)?.deprecatedAt || 'Unknown', + index: keyObj.__index })), endpoints: Object.entries(this.endpointStats) .filter(([_, stats]) => stats.success > 0 || stats.fail > 0) diff --git a/routes.js b/routes.js index 0818b6b..2fab60e 100644 --- a/routes.js +++ b/routes.js @@ -9,6 +9,7 @@ import { AnthropicResponseTransformer } from './transformers/response-anthropic. import { OpenAIResponseTransformer } from './transformers/response-openai.js'; import { getApiKey, recordRequestResult } from './auth.js'; import { getKeyManager } from './key-manager.js'; +import { isServerKeySet, setServerKey } from './server-auth.js'; const router = express.Router(); @@ -103,9 +104,12 @@ async function handleChatCompletions(req, res) { logInfo(`Routing to ${model.type} endpoint: ${endpoint.base_url}`); // Get API key (will auto-refresh if needed) + // When server access key is enabled, Authorization header is reserved for server auth. + // To pass a client-provided upstream token (client mode), use X-Endpoint-Authorization. let authHeader; try { - authHeader = await getApiKey(req.headers.authorization); + const clientUpstreamAuth = req.headers['x-endpoint-authorization'] || null; + authHeader = await getApiKey(clientUpstreamAuth); } catch (error) { logError('Failed to get API key', error); return res.status(500).json({ @@ -268,7 +272,8 @@ async function handleDirectResponses(req, res) { const clientAuthFromXApiKey = req.headers['x-api-key'] ? `Bearer ${req.headers['x-api-key']}` : null; - authHeader = await getApiKey(req.headers.authorization || clientAuthFromXApiKey); + const clientUpstreamAuth = req.headers['x-endpoint-authorization'] || null; + authHeader = await getApiKey(clientUpstreamAuth || clientAuthFromXApiKey); } catch (error) { logError('Failed to get API key', error); return res.status(500).json({ @@ -407,7 +412,8 @@ async function handleDirectMessages(req, res) { const clientAuthFromXApiKey = req.headers['x-api-key'] ? `Bearer ${req.headers['x-api-key']}` : null; - authHeader = await getApiKey(req.headers.authorization || clientAuthFromXApiKey); + const clientUpstreamAuth = req.headers['x-endpoint-authorization'] || null; + authHeader = await getApiKey(clientUpstreamAuth || clientAuthFromXApiKey); } catch (error) { logError('Failed to get API key', error); return res.status(500).json({ @@ -523,6 +529,45 @@ router.get('/status', (req, res) => { logInfo('GET /status'); try { + // First-visit setup: if no server key yet, present setup form + if (!isServerKeySet()) { + return res.send(` + + + + + droid2api Status - Setup + + + +

droid2api v2.0.0 Status

+
+
+

First-time setup: Set a server access key. After saving, all endpoints require this key via header Authorization: Bearer <key>. The /status page remains open.

+
+
+ + +
Remember this key. It will be stored locally on the server.
+ +
+
+ + + `); + } + const keyManager = getKeyManager(); if (!keyManager) { @@ -536,7 +581,7 @@ router.get('/status', (req, res) => { + + +
+

Server key saved successfully.

+

Use header Authorization: Bearer <key> for all API requests (except /status).

+

Go to Status

+
+ + + `); + } catch (error) { + logError('Error in POST /status/set-key', error); + res.status(500).send('Failed to set key'); + } +}); + // 注册路由 router.post('/v1/chat/completions', handleChatCompletions); router.post('/v1/responses', handleDirectResponses); router.post('/v1/messages', handleDirectMessages); export default router; + +// ---------------------- +// Balance query endpoints +// ---------------------- +// Helper to fetch usage/balance for a single key +async function fetchUsageForKeyRaw(apiKey) { + try { + const url = 'https://app.factory.ai/api/organization/members/chat-usage'; + const configUA = (() => { try { return getConfig()?.user_agent; } catch { return null; } })(); + const headers = { + 'Authorization': `Bearer ${apiKey}`, + 'User-Agent': configUA || 'factory-cli/0.19.3' + }; + const resp = await fetch(url, { method: 'GET', headers }); + if (!resp.ok) { + return { error: `HTTP ${resp.status}` }; + } + const data = await resp.json(); + const usage = data?.usage; + const standard = usage?.standard; + if (!standard) { + return { error: 'Invalid API response structure' }; + } + const totalAllowance = Number(standard.totalAllowance || 0); + const used = Number(standard.orgTotalTokensUsed || 0); + const usedRatio = typeof standard.usedRatio === 'number' ? standard.usedRatio : (totalAllowance > 0 ? used / totalAllowance : 0); + const startDate = usage?.startDate ? new Date(usage.startDate).toISOString() : null; + const endDate = usage?.endDate ? new Date(usage.endDate).toISOString() : null; + return { + totalAllowance, + used, + usedRatio, + startDate, + endDate + }; + } catch (e) { + return { error: 'Failed to fetch' }; + } +} + +// GET /status/balance/:index - fetch balance for specific key by index +router.get('/status/balance/:index', async (req, res) => { + try { + const keyManager = getKeyManager(); + if (!keyManager) { + return res.status(400).json({ error: 'Multi-key not configured' }); + } + const idx = parseInt(req.params.index, 10); + if (Number.isNaN(idx) || idx < 0 || idx >= keyManager.keys.length) { + return res.status(400).json({ error: 'Invalid index' }); + } + const keyObj = keyManager.keys[idx]; + if (!keyObj || keyObj.deprecated) { + return res.status(404).json({ error: 'Key not active' }); + } + const usage = await fetchUsageForKeyRaw(keyObj.key); + if (usage.error) { + return res.status(200).json({ index: idx, maskedKey: keyManager.maskKey(keyObj.key), error: usage.error, fetchedAt: new Date().toISOString() }); + } + const total = Number(usage.totalAllowance||0); + const used = Number(usage.used||0); + const remaining = Math.max(0, total - used); + // 同步到 KeyManager，余额用尽则在轮询中跳过 + try { keyManager.setBalanceByIndex(idx, { totalAllowance: total, used, remaining, fetchedAt: new Date().toISOString() }); } catch {} + return res.json({ index: idx, maskedKey: keyManager.maskKey(keyObj.key), ...usage, fetchedAt: new Date().toISOString() }); + } catch (error) { + logError('Error in GET /status/balance/:index', error); + return res.status(500).json({ error: 'Internal server error' }); + } +}); + +// GET /status/balances - fetch balances for all active keys +router.get('/status/balances', async (req, res) => { + try { + const keyManager = getKeyManager(); + if (!keyManager) { + return res.status(400).json({ error: 'Multi-key not configured' }); + } + const tasks = keyManager.keys + .map((k, i) => ({ k, i })) + .filter(x => !x.k.deprecated) + .map(async ({ k, i }) => { + const usage = await fetchUsageForKeyRaw(k.key); + if (usage.error) { + return { index: i, maskedKey: keyManager.maskKey(k.key), error: usage.error, fetchedAt: new Date().toISOString() }; + } + const total = Number(usage.totalAllowance||0); + const used = Number(usage.used||0); + const remaining = Math.max(0, total - used); + try { keyManager.setBalanceByIndex(i, { totalAllowance: total, used, remaining, fetchedAt: new Date().toISOString() }); } catch {} + return { index: i, maskedKey: keyManager.maskKey(k.key), ...usage, fetchedAt: new Date().toISOString() }; + }); + const list = await Promise.all(tasks); + return res.json(list); + } catch (error) { + logError('Error in GET /status/balances', error); + return res.status(500).json({ error: 'Internal server error' }); + } +}); + +// Skip-threshold settings (in-memory) +router.get('/status/skip-threshold', (req, res) => { + try { + const km = getKeyManager(); + const value = km ? km.getSkipThreshold?.() ?? 0 : 0; + res.json({ skipThreshold: value }); + } catch (e) { + res.status(200).json({ skipThreshold: 0 }); + } +}); + +router.post('/status/skip-threshold', express.json(), (req, res) => { + try { + const km = getKeyManager(); + if (!km) return res.status(400).json({ error: 'Multi-key not configured' }); + const body = req.body || {}; + const n = body.skipThreshold ?? body.n ?? body.remaining_threshold ?? 0; + km.setSkipThreshold?.(Number(n) || 0); + res.json({ ok: true, skipThreshold: km.getSkipThreshold?.() ?? 0 }); + } catch (e) { + res.status(500).json({ error: 'Failed to set threshold' }); + } +}); diff --git a/server-auth.js b/server-auth.js new file mode 100644 index 0000000..2a39abf --- /dev/null +++ b/server-auth.js @@ -0,0 +1,92 @@ +import fs from 'fs'; +import path from 'path'; +import { fileURLToPath } from 'url'; +import { logInfo, logError } from './logger.js'; + +const __filename = fileURLToPath(import.meta.url); +const __dirname = path.dirname(__filename); + +const KEY_FILE = path.join(__dirname, 'server-key.json'); + +let cachedKey = null; + +function loadKeyFromDisk() { + try { + if (fs.existsSync(KEY_FILE)) { + const raw = fs.readFileSync(KEY_FILE, 'utf-8'); + const data = JSON.parse(raw); + if (data && typeof data.key === 'string' && data.key.trim() !== '') { + cachedKey = data.key.trim(); + return cachedKey; + } + } + } catch (err) { + logError('Failed to read server key from disk', err); + } + return null; +} + +export function isServerKeySet() { + if (cachedKey && cachedKey.length > 0) return true; + const key = loadKeyFromDisk(); + return !!(key && key.length > 0); +} + +export function setServerKey(key) { + if (isServerKeySet()) { + throw new Error('Server key already set'); + } + if (!key || typeof key !== 'string' || key.trim() === '') { + throw new Error('Invalid key'); + } + const normalized = key.trim(); + try { + fs.writeFileSync(KEY_FILE, JSON.stringify({ key: normalized }, null, 2), 'utf-8'); + cachedKey = normalized; + logInfo('Server key has been set successfully'); + } catch (err) { + logError('Failed to write server key to disk', err); + throw err; + } +} + +export function verifyServerKey(provided) { + if (!provided || typeof provided !== 'string') return false; + const expected = cachedKey || loadKeyFromDisk(); + if (!expected) return false; + return expected === provided.trim(); +} + +export function serverAuthMiddleware(req, res, next) { + // Allow status and its subpaths without key + const path = req.path || req.originalUrl || ''; + if (path === '/status' || path.startsWith('/status/')) { + return next(); + } + + // If key is not set yet, block all other routes and instruct to visit /status + if (!isServerKeySet()) { + return res.status(503).json({ + error: 'Server key not set', + message: 'Visit /status to set the initial access key.' + }); + } + + // Accept key via Authorization: Bearer + const authHeader = req.headers['authorization']; + let provided = null; + if (typeof authHeader === 'string') { + const parts = authHeader.split(' '); + if (parts.length === 2 && /^Bearer$/i.test(parts[0])) { + provided = parts[1]; + } + } + if (!verifyServerKey(typeof provided === 'string' ? provided : '')) { + return res.status(401).json({ + error: 'Unauthorized', + message: 'Missing or invalid Authorization: Bearer ' + }); + } + + return next(); +} diff --git a/server.js b/server.js index 90e84d2..ed0516d 100644 --- a/server.js +++ b/server.js @@ -2,6 +2,7 @@ import express from 'express'; import { loadConfig, isDevMode, getPort, getRoundRobin, getRemoveOn402 } from './config.js'; import { logInfo, logError } from './logger.js'; import router from './routes.js'; +import { serverAuthMiddleware } from './server-auth.js'; import { initializeAuth } from './auth.js'; const app = express(); @@ -12,7 +13,7 @@ app.use(express.urlencoded({ extended: true, limit: '50mb' })); app.use((req, res, next) => { res.header('Access-Control-Allow-Origin', '*'); res.header('Access-Control-Allow-Methods', 'GET, POST, PUT, DELETE, OPTIONS'); - res.header('Access-Control-Allow-Headers', 'Content-Type, Authorization, X-API-Key, anthropic-version'); + res.header('Access-Control-Allow-Headers', 'Content-Type, Authorization, X-API-Key, anthropic-version, X-Endpoint-Authorization'); if (req.method === 'OPTIONS') { return res.sendStatus(200); @@ -20,6 +21,9 @@ app.use((req, res, next) => { next(); }); +// Enforce server-wide access key after first-time setup via /status +app.use(serverAuthMiddleware); + app.use(router); app.get('/', (req, res) => {