Add per-user daily token quota with persistent usage tracking

This commit is contained in:
2026-08-21 12:25:57 -07:00
parent 2c3dc66878
commit f1200b7d79
6 changed files with 199 additions and 45 deletions

157
server.js
View File

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