Add per-user daily token quota with persistent usage tracking
This commit is contained in:
157
server.js
157
server.js
@@ -3,6 +3,7 @@
|
||||
* Serves the static tool files.
|
||||
* Proxies AI draft requests to the Anthropic API, so the API key
|
||||
* stays on the server and never appears in the browser.
|
||||
* Identifies each caller by login and enforces a daily token quota per person.
|
||||
*/
|
||||
const http = require('http');
|
||||
const fs = require('fs');
|
||||
@@ -49,17 +50,52 @@ function sendJson(res, status, obj) {
|
||||
}
|
||||
|
||||
/* ===== Access control =====
|
||||
* Set APP_USERNAME and APP_PASSWORD (in .env or the container environment)
|
||||
* to require a login for the whole app. Leave both unset for solo local use
|
||||
* with no login prompt. This is a simple shared-credential gate, meant to sit
|
||||
* behind the company VPN and internal network, not to replace them. */
|
||||
* Named users. Set APP_USERS as a comma-separated list of user:password
|
||||
* pairs, for example: APP_USERS=alice:pass1,bob:pass2
|
||||
* Each name is also the identity used for the daily token quota below.
|
||||
*
|
||||
* Legacy single-user mode: set APP_USERNAME and APP_PASSWORD instead. That
|
||||
* name becomes the one identity everyone shares (no per-person quota).
|
||||
*
|
||||
* Leave all of the above unset for solo local use: no login prompt, and
|
||||
* usage is tracked under the identity "local".
|
||||
*
|
||||
* This is a simple shared-credential gate, meant to sit behind the company
|
||||
* VPN and internal network, not to replace them. */
|
||||
function loadUserMap() {
|
||||
const map = {};
|
||||
if (process.env.APP_USERS) {
|
||||
process.env.APP_USERS.split(',').forEach((pair) => {
|
||||
const idx = pair.indexOf(':');
|
||||
if (idx === -1) return;
|
||||
const name = pair.slice(0, idx).trim();
|
||||
const pass = pair.slice(idx + 1).trim();
|
||||
if (name && pass) map[name] = pass;
|
||||
});
|
||||
} else if (process.env.APP_USERNAME && process.env.APP_PASSWORD) {
|
||||
map[process.env.APP_USERNAME] = process.env.APP_PASSWORD;
|
||||
}
|
||||
return map;
|
||||
}
|
||||
const USER_MAP = loadUserMap();
|
||||
const AUTH_REQUIRED = Object.keys(USER_MAP).length > 0;
|
||||
|
||||
// Returns the identified username on success, or null on failure.
|
||||
function checkBasicAuth(req) {
|
||||
const user = process.env.APP_USERNAME;
|
||||
const pass = process.env.APP_PASSWORD;
|
||||
if (!user || !pass) return true;
|
||||
if (!AUTH_REQUIRED) return 'local';
|
||||
const header = req.headers['authorization'] || '';
|
||||
const expected = 'Basic ' + Buffer.from(`${user}:${pass}`).toString('base64');
|
||||
return header === expected;
|
||||
if (!header.startsWith('Basic ')) return null;
|
||||
let decoded;
|
||||
try {
|
||||
decoded = Buffer.from(header.slice(6), 'base64').toString('utf8');
|
||||
} catch (err) {
|
||||
return null;
|
||||
}
|
||||
const idx = decoded.indexOf(':');
|
||||
if (idx === -1) return null;
|
||||
const name = decoded.slice(0, idx);
|
||||
const pass = decoded.slice(idx + 1);
|
||||
return USER_MAP[name] === pass ? name : null;
|
||||
}
|
||||
function requireAuth(res) {
|
||||
res.writeHead(401, {
|
||||
@@ -70,9 +106,9 @@ function requireAuth(res) {
|
||||
}
|
||||
|
||||
/* ===== Rate limiting for the AI proxy =====
|
||||
* A shared company API key means one leaked link or one runaway script can
|
||||
* generate real cost. This caps AI draft calls per source IP. It resets if
|
||||
* the process restarts; that is an accepted tradeoff for a small internal tool. */
|
||||
* Caps AI draft calls per source IP, independent of the token quota below.
|
||||
* This catches a runaway script quickly, before it burns through a whole
|
||||
* day's token quota in a few seconds. Resets if the process restarts. */
|
||||
const RATE_LIMIT_MAX = parseInt(process.env.RATE_LIMIT_MAX || '20', 10);
|
||||
const RATE_LIMIT_WINDOW_MS = parseInt(process.env.RATE_LIMIT_WINDOW_MS || String(5 * 60 * 1000), 10);
|
||||
const rateLimitHits = new Map();
|
||||
@@ -84,13 +120,66 @@ function isRateLimited(ip) {
|
||||
return hits.length > RATE_LIMIT_MAX;
|
||||
}
|
||||
|
||||
/* ===== Per-user daily token quota =====
|
||||
* TOKEN_LIMIT_PER_USER caps combined input+output tokens per identified user,
|
||||
* per UTC calendar day. Usage is written to TOKEN_USAGE_FILE after every AI
|
||||
* call, so it survives a container restart. Mount that file's folder as a
|
||||
* volume in Docker or it resets on every redeploy.
|
||||
*
|
||||
* The check runs before the API call using the day's tally so far: if the
|
||||
* user is already at or over the limit, the call is refused with no cost.
|
||||
* A single call in progress when the limit is reached can still push the
|
||||
* tally past the cap by that one call's tokens; the next call is blocked.
|
||||
* That is an accepted tradeoff, since real token cost of a call is only
|
||||
* known once its response returns. */
|
||||
const TOKEN_LIMIT_PER_USER = parseInt(process.env.TOKEN_LIMIT_PER_USER || '50000', 10);
|
||||
const TOKEN_USAGE_FILE = process.env.TOKEN_USAGE_FILE || path.join(ROOT, 'data', 'token-usage.json');
|
||||
|
||||
function todayUTC() {
|
||||
return new Date().toISOString().slice(0, 10); // "YYYY-MM-DD"
|
||||
}
|
||||
function loadUsage() {
|
||||
try {
|
||||
return JSON.parse(fs.readFileSync(TOKEN_USAGE_FILE, 'utf8'));
|
||||
} catch (err) {
|
||||
return {};
|
||||
}
|
||||
}
|
||||
function saveUsage(usage) {
|
||||
try {
|
||||
fs.mkdirSync(path.dirname(TOKEN_USAGE_FILE), { recursive: true });
|
||||
fs.writeFileSync(TOKEN_USAGE_FILE, JSON.stringify(usage));
|
||||
} catch (err) {
|
||||
console.error('Could not write token usage file:', err.message);
|
||||
}
|
||||
}
|
||||
function usedTokensToday(usage, user) {
|
||||
const rec = usage[user];
|
||||
if (!rec || rec.date !== todayUTC()) return 0;
|
||||
return rec.tokens;
|
||||
}
|
||||
function addTokens(user, tokens) {
|
||||
const usage = loadUsage();
|
||||
const today = todayUTC();
|
||||
const rec = usage[user] && usage[user].date === today ? usage[user] : { date: today, tokens: 0 };
|
||||
rec.tokens += tokens;
|
||||
usage[user] = rec;
|
||||
saveUsage(usage);
|
||||
return rec.tokens;
|
||||
}
|
||||
function secondsUntilUTCMidnight() {
|
||||
const now = new Date();
|
||||
const midnight = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), now.getUTCDate() + 1));
|
||||
return Math.round((midnight - now) / 1000);
|
||||
}
|
||||
|
||||
/* ===== Mock mode =====
|
||||
* Set MOCK_AI=true to test the whole app, including the AI draft button,
|
||||
* with no API key and no real API call. Returns a canned breadcrumb draft
|
||||
* shaped like a real Anthropic response, so the frontend parsing and gate
|
||||
* checks run exactly as they would against the live API. Use this for
|
||||
* testing the login gate, rate limit, and Docker setup without cost. */
|
||||
function mockClaudeResponse(res) {
|
||||
* Set MOCK_AI=true to test the whole app, including the AI draft button and
|
||||
* the token quota above, with no API key and no real API call. Returns a
|
||||
* canned breadcrumb draft with a synthetic usage count, shaped like a real
|
||||
* Anthropic response, so the frontend parsing, gate checks, and quota
|
||||
* accounting all run exactly as they would against the live API. */
|
||||
function mockClaudeResponse(res, user) {
|
||||
const draft = {
|
||||
domain: 'MOCK: Progress Visibility',
|
||||
driver: 'MOCK driver: schedule forecasts drift from field reality.',
|
||||
@@ -99,16 +188,19 @@ function mockClaudeResponse(res) {
|
||||
method: 'MOCK method: monthly spot audit of 30+ sampled assets.',
|
||||
baseline: 'MOCK baseline: first audit result.'
|
||||
};
|
||||
const usage = { input_tokens: 950, output_tokens: 280 };
|
||||
const body = {
|
||||
content: [{ type: 'text', text: JSON.stringify(draft) }]
|
||||
content: [{ type: 'text', text: JSON.stringify(draft) }],
|
||||
usage
|
||||
};
|
||||
addTokens(user, usage.input_tokens + usage.output_tokens);
|
||||
res.writeHead(200, { 'Content-Type': 'application/json; charset=utf-8' });
|
||||
res.end(JSON.stringify(body));
|
||||
}
|
||||
|
||||
function proxyToClaude(req, res) {
|
||||
function proxyToClaude(req, res, user) {
|
||||
if (String(process.env.MOCK_AI).toLowerCase() === 'true') {
|
||||
mockClaudeResponse(res);
|
||||
mockClaudeResponse(res, user);
|
||||
return;
|
||||
}
|
||||
if (!process.env.ANTHROPIC_API_KEY) {
|
||||
@@ -133,6 +225,14 @@ function proxyToClaude(req, res) {
|
||||
const text = await upstream.text();
|
||||
res.writeHead(upstream.status, { 'Content-Type': 'application/json; charset=utf-8' });
|
||||
res.end(text);
|
||||
try {
|
||||
const parsed = JSON.parse(text);
|
||||
if (parsed.usage) {
|
||||
addTokens(user, (parsed.usage.input_tokens || 0) + (parsed.usage.output_tokens || 0));
|
||||
}
|
||||
} catch (err) {
|
||||
// Response was not JSON, or had no usage field. Nothing to record.
|
||||
}
|
||||
} catch (err) {
|
||||
sendJson(res, 502, { error: 'Could not reach the Anthropic API: ' + err.message });
|
||||
}
|
||||
@@ -172,7 +272,8 @@ function serveStatic(req, res) {
|
||||
}
|
||||
|
||||
const server = http.createServer((req, res) => {
|
||||
if (!checkBasicAuth(req)) {
|
||||
const user = checkBasicAuth(req);
|
||||
if (!user) {
|
||||
requireAuth(res);
|
||||
return;
|
||||
}
|
||||
@@ -182,7 +283,14 @@ const server = http.createServer((req, res) => {
|
||||
sendJson(res, 429, { error: `Rate limit reached (${RATE_LIMIT_MAX} AI drafts per ${Math.round(RATE_LIMIT_WINDOW_MS / 60000)} min). Wait a bit and try again.` });
|
||||
return;
|
||||
}
|
||||
proxyToClaude(req, res);
|
||||
const usedToday = usedTokensToday(loadUsage(), user);
|
||||
if (usedToday >= TOKEN_LIMIT_PER_USER) {
|
||||
sendJson(res, 429, {
|
||||
error: `Daily token quota reached (${usedToday}/${TOKEN_LIMIT_PER_USER} tokens for "${user}"). Resets in ${secondsUntilUTCMidnight()}s, at UTC midnight.`
|
||||
});
|
||||
return;
|
||||
}
|
||||
proxyToClaude(req, res, user);
|
||||
return;
|
||||
}
|
||||
if (req.method === 'GET') {
|
||||
@@ -200,5 +308,6 @@ server.listen(PORT, () => {
|
||||
} else {
|
||||
console.log(process.env.ANTHROPIC_API_KEY ? 'ANTHROPIC_API_KEY loaded: AI draft button is live.' : 'No ANTHROPIC_API_KEY found: AI draft button will return an error until you add one to .env.');
|
||||
}
|
||||
console.log((process.env.APP_USERNAME && process.env.APP_PASSWORD) ? 'Login required: APP_USERNAME/APP_PASSWORD are set.' : 'No login required: APP_USERNAME/APP_PASSWORD are not set.');
|
||||
console.log(AUTH_REQUIRED ? `Login required: ${Object.keys(USER_MAP).length} named user(s) configured.` : 'No login required: no APP_USERS or APP_USERNAME/APP_PASSWORD are set.');
|
||||
console.log(`Per-user daily token quota: ${TOKEN_LIMIT_PER_USER} tokens. Usage file: ${TOKEN_USAGE_FILE}`);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user