token-consumer / server.js
qwen2api's picture
Upload 11 files
8a81249 verified
Raw
History Blame Contribute Delete
25.9 kB
const express = require('express');
const cors = require('cors');
const { v4: uuidv4 } = require('uuid');
const path = require('path');
const os = require('os');
const storage = require('./storage');
const { TaskExecutor, HttpTaskExecutor, fetchModels } = require('./executor');
const log = require('./logger');
const app = express();
const PORT = process.env.PORT || 51730;
// Auth configuration
const AUTH_USER = process.env.AUTH_USER || 'admin';
const AUTH_PASS = process.env.AUTH_PASS || 'admin';
// Rate limiter - simple in-memory implementation
const rateLimitMap = new Map();
const RATE_LIMIT_WINDOW = 60000; // 1 minute
const RATE_LIMIT_MAX = 300; // max requests per window
function rateLimiter(req, res, next) {
const ip = req.ip || req.connection.remoteAddress || 'unknown';
const now = Date.now();
let entry = rateLimitMap.get(ip);
if (!entry || now - entry.startTime > RATE_LIMIT_WINDOW) {
entry = { count: 1, startTime: now };
rateLimitMap.set(ip, entry);
} else {
entry.count++;
}
if (entry.count > RATE_LIMIT_MAX) {
return res.status(429).json({ error: '请求过于频繁,请稍后再试' });
}
// Clean up old entries periodically
if (rateLimitMap.size > 1000) {
for (const [key, val] of rateLimitMap) {
if (now - val.startTime > RATE_LIMIT_WINDOW) {
rateLimitMap.delete(key);
}
}
}
next();
}
// Middleware
app.use(cors());
app.use(express.json());
app.use(rateLimiter);
// HTTP Basic Auth middleware for /api routes
const basicAuth = (req, res, next) => {
const authHeader = req.headers.authorization;
if (!authHeader || !authHeader.startsWith('Basic ')) {
res.setHeader('WWW-Authenticate', 'Basic realm="Token Consumer"');
return res.status(401).json({ error: '需要认证' });
}
const base64Credentials = authHeader.split(' ')[1];
const credentials = Buffer.from(base64Credentials, 'base64').toString('utf8');
const [username, password] = credentials.split(':');
if (username === AUTH_USER && password === AUTH_PASS) {
return next();
}
res.setHeader('WWW-Authenticate', 'Basic realm="Token Consumer"');
return res.status(401).json({ error: '认证失败' });
};
// Apply auth to all /api routes
app.use('/api', basicAuth);
// Health check endpoint (no auth required)
app.get('/health', (req, res) => {
res.json({
status: 'ok',
timestamp: new Date().toISOString(),
uptime: process.uptime()
});
});
// Serve static files
app.use(express.static(__dirname));
// In-memory task executors
const executors = new Map();
// Create or update task
app.post('/api/tasks', (req, res) => {
try {
const { id, label, config } = req.body;
log.debug('POST /api/tasks', { id, label, model: config?.model });
if (!config?.base || !config?.token || !config?.model || !config?.usr) {
log.warn('缺少必要参数');
return res.status(400).json({ error: '缺少必要参数' });
}
const taskId = id || uuidv4();
const now = Date.now();
const task = {
id: taskId,
label: label || `任务-${taskId.slice(0, 6)}`,
config: {
base: config.base,
token: config.token,
model: config.model,
sys: config.sys || '',
usr: config.usr,
loop: Math.max(1, Math.min(10000, Number(config.loop) || 10)),
threads: Math.max(1, Math.min(200, Number(config.threads) || 3)),
max: Math.max(1, Math.min(32768, Number(config.max) || 1024)),
temp: Math.max(0, Math.min(2, Number(config.temp) || 1)),
timeout: Math.max(5000, Number(config.timeout) || 60000),
maxTokens: Math.max(1000, Number(config.maxTokens) || 1000000000),
randOn: !!config.randOn,
streamOn: config.streamOn !== false,
ratioTarget: Math.max(0.1, Math.min(20, Number(config.ratioTarget) || 1)),
markerLen: Math.max(0, Number(config.markerLen) || 24),
waitBetween: Math.max(0, Number(config.waitBetween) || 1000)
},
status: 'idle',
stats: {
total: Math.max(1, Math.min(10000, Number(config.loop) || 10)) * Math.max(1, Math.min(200, Number(config.threads) || 3)),
completed: 0,
success: 0,
failed: 0,
aborted: 0,
promptTokens: 0,
completionTokens: 0,
totalTokens: 0
},
createdAt: storage.getTask(taskId)?.createdAt || now,
updatedAt: now
};
storage.saveTask(taskId, task);
log.info('任务已保存:', taskId, label || task.label);
res.json(task);
} catch (e) {
log.error('保存任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Get all tasks
app.get('/api/tasks', (req, res) => {
try {
const tasks = storage.getAllTasks();
const result = Object.values(tasks).map(t => {
const executor = executors.get(t.id);
// Fix stats.total for old tasks
if (!t.stats.total && t.config) {
t.stats.total = (t.config.loop || 10) * (t.config.threads || 3);
}
return {
...t,
running: executor?.running || false,
currentStats: executor?.getStatus() || null
};
});
log.debug('GET /api/tasks, count:', result.length);
res.json(result);
} catch (e) {
log.error('获取任务列表失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Get single task
app.get('/api/tasks/:id', (req, res) => {
try {
const task = storage.getTask(req.params.id);
if (!task) {
return res.status(404).json({ error: '任务不存在' });
}
const executor = executors.get(req.params.id);
// Fix stats.total for old tasks
if (!task.stats.total && task.config) {
task.stats.total = (task.config.loop || 10) * (task.config.threads || 3);
}
res.json({
...task,
running: executor?.running || false,
currentStats: executor?.getStatus() || null
});
} catch (e) {
res.status(500).json({ error: e.message });
}
});
// Delete task
app.delete('/api/tasks/:id', (req, res) => {
try {
log.info('删除任务:', req.params.id);
const executor = executors.get(req.params.id);
if (executor?.running) {
log.warn('任务正在运行,先停止:', req.params.id);
executor.stop();
}
executors.delete(req.params.id);
storage.deleteTask(req.params.id);
log.info('任务已删除:', req.params.id);
res.json({ success: true });
} catch (e) {
log.error('删除任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Start task
app.post('/api/tasks/:id/start', (req, res) => {
try {
const task = storage.getTask(req.params.id);
if (!task) {
log.warn('任务不存在:', req.params.id);
return res.status(404).json({ error: '任务不存在' });
}
let executor = executors.get(req.params.id);
if (executor?.running && !executor.paused) {
log.warn('任务已在运行:', req.params.id);
return res.status(400).json({ error: '任务已在运行' });
}
// If paused, resume
if (executor?.paused) {
log.info('恢复暂停的任务:', req.params.id);
executor.resume();
return res.json({ success: true, message: '任务已恢复' });
}
log.info('启动任务:', req.params.id, task.label);
// Check if we have saved progress to resume
const savedProgress = task.progress || null;
if (savedProgress) {
log.info('从保存的进度恢复:', req.params.id);
}
// Create new executor
executor = new TaskExecutor(req.params.id, task.config, (status) => {
// Update task in storage
const updatedTask = storage.getTask(req.params.id);
if (updatedTask) {
updatedTask.stats = status.stats;
updatedTask.status = status.running ? (status.paused ? 'paused' : 'running') : (status.stopped ? 'stopped' : 'completed');
// Always save progress with threadLogs for display, even after completion
updatedTask.progress = status;
updatedTask.updatedAt = Date.now();
storage.saveTask(req.params.id, updatedTask);
}
// Remove from memory if stopped/completed
if (!status.running) {
executors.delete(req.params.id);
log.info('任务结束:', req.params.id, '总tokens:', status.stats?.totalTokens || 0);
}
}, savedProgress);
executors.set(req.params.id, executor);
log.debug('Executor created and set in map:', req.params.id, 'map size:', executors.size);
// Start in background
executor.start().then(() => {
log.debug('Executor start() resolved:', req.params.id);
}).catch(e => log.error('任务执行错误:', e.message));
res.json({ success: true, message: savedProgress ? '任务已恢复' : '任务已启动' });
} catch (e) {
log.error('启动任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Pause task
app.post('/api/tasks/:id/pause', (req, res) => {
try {
const executor = executors.get(req.params.id);
if (!executor || !executor.running || executor.paused) {
log.warn('任务未在运行或已暂停:', req.params.id);
return res.status(400).json({ error: '任务未在运行或已暂停' });
}
log.info('暂停任务:', req.params.id);
executor.pause();
res.json({ success: true, message: '任务已暂停' });
} catch (e) {
log.error('暂停任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Stop task (reset progress)
app.post('/api/tasks/:id/stop', (req, res) => {
try {
const executor = executors.get(req.params.id);
if (!executor || !executor.running) {
log.warn('任务未在运行:', req.params.id);
return res.status(400).json({ error: '任务未在运行' });
}
log.info('停止任务:', req.params.id);
executor.stop();
// Update task status (progress will be saved by executor callback)
const task = storage.getTask(req.params.id);
if (task) {
task.status = 'stopped';
storage.saveTask(req.params.id, task);
}
res.json({ success: true, message: '任务已停止' });
} catch (e) {
log.error('停止任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Get task status (for polling)
app.get('/api/tasks/:id/status', (req, res) => {
try {
const executor = executors.get(req.params.id);
const task = storage.getTask(req.params.id);
if (!task) {
return res.status(404).json({ error: '任务不存在' });
}
// Use getStatus() for running tasks, or progress from storage for completed tasks
const status = executor?.getStatus();
const progress = status || task.progress;
res.json({
taskId: req.params.id,
running: status?.running || false,
stopped: status?.stopped || task.status === 'stopped' || false,
paused: status?.paused || false,
stats: progress?.stats || task.stats,
ratio: progress?.ratio || null,
elapsed: progress?.elapsed || 0,
threadLogs: progress?.threadLogs || {}
});
} catch (e) {
log.error('获取任务状态失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Load models
app.post('/api/models', async (req, res) => {
try {
const { base, token } = req.body;
log.debug('加载模型列表, base:', base);
if (!base || !token) {
log.warn('缺少 base 或 token');
return res.status(400).json({ error: '缺少 base 或 token' });
}
const models = await fetchModels(base, token);
log.info('模型列表加载成功, 数量:', models.length);
res.json({ models });
} catch (e) {
log.error('加载模型失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Save config to server
app.post('/api/config', (req, res) => {
try {
storage.saveConfig(req.body);
res.json({ success: true });
} catch (e) {
res.status(500).json({ error: e.message });
}
});
// Get config from server
app.get('/api/config', (req, res) => {
try {
const config = storage.loadConfig();
res.json(config);
} catch (e) {
res.status(500).json({ error: e.message });
}
});
// Get memory usage
app.get('/api/memory', (req, res) => {
try {
// Process memory
const processMem = process.memoryUsage();
// System memory
const totalMem = os.totalmem();
const freeMem = os.freemem();
const usedMem = totalMem - freeMem;
res.json({
process: {
rss: processMem.rss, // Resident Set Size
heapTotal: processMem.heapTotal,
heapUsed: processMem.heapUsed,
external: processMem.external,
arrayBuffers: processMem.arrayBuffers
},
system: {
total: totalMem,
free: freeMem,
used: usedMem,
usagePercent: Math.round(usedMem / totalMem * 100)
}
});
} catch (e) {
res.status(500).json({ error: e.message });
}
});
// Get raw tasks.json content
app.get('/api/tasks-raw', (req, res) => {
try {
const data = storage.getTasksRaw();
res.json({ content: JSON.stringify(data, null, 2) });
} catch (e) {
log.error('读取 tasks.json 失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Save raw tasks.json content
app.post('/api/tasks-raw', async (req, res) => {
try {
// Check if any task is running (not paused - paused tasks don't write)
const runningCount = [...executors.values()].filter(e => e.running && !e.paused).length;
if (runningCount > 0) {
return res.status(409).json({
error: `有 ${runningCount} 个任务正在运行,请先暂停或停止所有运行中的任务后再保存配置`
});
}
const { content } = req.body;
if (!content) {
return res.status(400).json({ error: '内容不能为空' });
}
// Validate JSON
try {
JSON.parse(content);
} catch (e) {
return res.status(400).json({ error: 'JSON 格式无效: ' + e.message });
}
storage.saveTasksRaw(content);
log.info('tasks.json 已更新');
res.json({ success: true, message: '保存成功' });
} catch (e) {
log.error('保存 tasks.json 失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// HTTP Task endpoints
// Create or update HTTP task
app.post('/api/http-tasks', (req, res) => {
try {
const { id, label, config } = req.body;
log.debug('POST /api/http-tasks', { id, label });
if (!config?.url || !config?.method) {
log.warn('缺少必要参数');
return res.status(400).json({ error: '缺少必要参数(url和method)' });
}
const taskId = id || uuidv4();
const now = Date.now();
const task = {
id: taskId,
label: label || `HTTP任务-${taskId.slice(0, 6)}`,
type: 'http',
config: {
url: config.url,
method: config.method || 'GET',
headers: config.headers || '{}',
body: config.body || '',
loop: Math.max(1, Math.min(10000, Number(config.loop) || 10)),
threads: Math.max(1, Math.min(200, Number(config.threads) || 3)),
timeout: Math.max(5000, Number(config.timeout) || 60000),
waitBetween: Math.max(0, Number(config.waitBetween) || 1000)
},
status: 'idle',
stats: {
total: Math.max(1, Math.min(10000, Number(config.loop) || 10)) * Math.max(1, Math.min(200, Number(config.threads) || 3)),
completed: 0,
success: 0,
failed: 0,
aborted: 0,
totalRequests: 0
},
createdAt: storage.getTask(taskId)?.createdAt || now,
updatedAt: now
};
storage.saveTask(taskId, task);
log.info('HTTP任务已保存:', taskId, label || task.label);
res.json(task);
} catch (e) {
log.error('保存HTTP任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Get all HTTP tasks
app.get('/api/http-tasks', (req, res) => {
try {
const allTasks = storage.getAllTasks();
const tasks = Object.values(allTasks).filter(t => t.type === 'http').map(t => {
const executor = executors.get(t.id);
if (!t.stats.total && t.config) {
t.stats.total = (t.config.loop || 10) * (t.config.threads || 3);
}
return {
...t,
running: executor?.running || false,
currentStats: executor?.getStatus() || null
};
});
log.debug('GET /api/http-tasks, count:', tasks.length);
res.json(tasks);
} catch (e) {
log.error('获取HTTP任务列表失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Get single HTTP task
app.get('/api/http-tasks/:id', (req, res) => {
try {
const task = storage.getTask(req.params.id);
if (!task || task.type !== 'http') {
return res.status(404).json({ error: 'HTTP任务不存在' });
}
const executor = executors.get(req.params.id);
if (!task.stats.total && task.config) {
task.stats.total = (task.config.loop || 10) * (task.config.threads || 3);
}
res.json({
...task,
running: executor?.running || false,
currentStats: executor?.getStatus() || null
});
} catch (e) {
res.status(500).json({ error: e.message });
}
});
// Delete HTTP task
app.delete('/api/http-tasks/:id', (req, res) => {
try {
log.info('删除HTTP任务:', req.params.id);
const executor = executors.get(req.params.id);
if (executor?.running) {
log.warn('HTTP任务正在运行,先停止:', req.params.id);
executor.stop();
}
executors.delete(req.params.id);
storage.deleteTask(req.params.id);
log.info('HTTP任务已删除:', req.params.id);
res.json({ success: true });
} catch (e) {
log.error('删除HTTP任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Start HTTP task
app.post('/api/http-tasks/:id/start', (req, res) => {
try {
const task = storage.getTask(req.params.id);
if (!task || task.type !== 'http') {
log.warn('HTTP任务不存在:', req.params.id);
return res.status(404).json({ error: 'HTTP任务不存在' });
}
let executor = executors.get(req.params.id);
if (executor?.running && !executor.paused) {
log.warn('HTTP任务已在运行:', req.params.id);
return res.status(400).json({ error: 'HTTP任务已在运行' });
}
if (executor?.paused) {
log.info('恢复暂停的HTTP任务:', req.params.id);
executor.resume();
return res.json({ success: true, message: '任务已恢复' });
}
log.info('启动HTTP任务:', req.params.id, task.label);
const savedProgress = task.progress || null;
executor = new HttpTaskExecutor(req.params.id, task.config, (status) => {
const updatedTask = storage.getTask(req.params.id);
if (updatedTask) {
updatedTask.stats = status.stats;
updatedTask.status = status.running ? (status.paused ? 'paused' : 'running') : (status.stopped ? 'stopped' : 'completed');
updatedTask.progress = status;
updatedTask.updatedAt = Date.now();
storage.saveTask(req.params.id, updatedTask);
}
if (!status.running) {
executors.delete(req.params.id);
log.info('HTTP任务结束:', req.params.id);
}
}, savedProgress);
executors.set(req.params.id, executor);
executor.start().then(() => {
log.debug('HTTP Executor start() resolved:', req.params.id);
}).catch(e => log.error('HTTP任务执行错误:', e.message));
res.json({ success: true, message: savedProgress ? '任务已恢复' : '任务已启动' });
} catch (e) {
log.error('启动HTTP任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Pause HTTP task
app.post('/api/http-tasks/:id/pause', (req, res) => {
try {
const executor = executors.get(req.params.id);
if (!executor || !executor.running || executor.paused) {
log.warn('HTTP任务未在运行或已暂停:', req.params.id);
return res.status(400).json({ error: '任务未在运行或已暂停' });
}
log.info('暂停HTTP任务:', req.params.id);
executor.pause();
res.json({ success: true, message: '任务已暂停' });
} catch (e) {
log.error('暂停HTTP任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Stop HTTP task
app.post('/api/http-tasks/:id/stop', (req, res) => {
try {
const executor = executors.get(req.params.id);
if (!executor || !executor.running) {
log.warn('HTTP任务未在运行:', req.params.id);
return res.status(400).json({ error: '任务未在运行' });
}
log.info('停止HTTP任务:', req.params.id);
executor.stop();
const task = storage.getTask(req.params.id);
if (task) {
task.status = 'stopped';
storage.saveTask(req.params.id, task);
}
res.json({ success: true, message: '任务已停止' });
} catch (e) {
log.error('停止HTTP任务失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Get HTTP task status
app.get('/api/http-tasks/:id/status', (req, res) => {
try {
const executor = executors.get(req.params.id);
const task = storage.getTask(req.params.id);
if (!task || task.type !== 'http') {
return res.status(404).json({ error: 'HTTP任务不存在' });
}
const status = executor?.getStatus();
const progress = status || task.progress;
res.json({
taskId: req.params.id,
running: status?.running || false,
stopped: status?.stopped || task.status === 'stopped' || false,
paused: status?.paused || false,
stats: progress?.stats || task.stats,
elapsed: progress?.elapsed || 0,
threadLogs: progress?.threadLogs || {}
});
} catch (e) {
log.error('获取HTTP任务状态失败:', e.message);
res.status(500).json({ error: e.message });
}
});
// Serve index.html
app.get('/', (req, res) => {
res.sendFile(path.join(__dirname, 'token-consumer.html'));
});
// Resume running tasks on startup
function resumeRunningTasks() {
const tasks = storage.getAllTasks();
const runningTasks = Object.values(tasks).filter(t => t.status === 'running' || t.status === 'paused');
if (runningTasks.length === 0) {
log.info('没有需要恢复的运行中任务');
return;
}
log.info(`发现 ${runningTasks.length} 个运行中任务,正在恢复...`);
for (const task of runningTasks) {
try {
const savedProgress = task.progress || null;
const executor = new TaskExecutor(task.id, task.config, (status) => {
const updatedTask = storage.getTask(task.id);
if (updatedTask) {
updatedTask.stats = status.stats;
updatedTask.status = status.running ? (status.paused ? 'paused' : 'running') : (status.stopped ? 'stopped' : 'completed');
// Always save progress with threadLogs for display
updatedTask.progress = status;
updatedTask.updatedAt = Date.now();
storage.saveTask(task.id, updatedTask);
}
if (!status.running) {
executors.delete(task.id);
log.info('任务结束:', task.id, '总tokens:', status.stats?.totalTokens || 0);
}
}, savedProgress);
executors.set(task.id, executor);
executor.start().catch(e => log.error('恢复任务失败:', task.id, e.message));
log.info('任务已恢复:', task.id, task.label);
} catch (e) {
log.error('恢复任务失败:', task.id, e.message);
}
}
}
const HOST = process.env.HOST || '0.0.0.0';
app.listen(PORT, HOST, () => {
console.log(`Server running at http://${HOST}:${PORT}`);
console.log(`Log level: ${(process.env.LOG_LEVEL || 'info').toUpperCase()} (set LOG_LEVEL env: debug|info|warn|error|none)`);
console.log(`API endpoints:`);
console.log(` GET /health - 健康检查`);
console.log(` GET /api/tasks - 获取所有任务`);
console.log(` POST /api/tasks - 创建/更新任务`);
console.log(` GET /api/tasks/:id - 获取单个任务`);
console.log(` DELETE /api/tasks/:id - 删除任务`);
console.log(` POST /api/tasks/:id/start - 启动/恢复任务`);
console.log(` POST /api/tasks/:id/pause - 暂停任务`);
console.log(` POST /api/tasks/:id/stop - 停止任务`);
console.log(` GET /api/tasks/:id/status - 获取任务状态`);
console.log(` POST /api/models - 加载模型列表`);
console.log(` GET /api/config - 获取配置`);
console.log(` POST /api/config - 保存配置`);
console.log(` GET /api/memory - 获取内存使用`);
// Resume running tasks after server starts
resumeRunningTasks();
});
// Graceful shutdown handler
function gracefulShutdown(signal) {
console.log(`\n[${signal}] Graceful shutdown initiated...`);
// Stop all running executors
let runningCount = 0;
for (const [id, executor] of executors) {
if (executor.running) {
executor.stop();
runningCount++;
}
}
if (runningCount > 0) {
console.log(`[Shutdown] Stopped ${runningCount} running task(s)`);
}
// Flush any pending data
storage.shutdown();
console.log('[Shutdown] Complete');
process.exit(0);
}
// Handle termination signals
process.on('SIGINT', () => gracefulShutdown('SIGINT'));
process.on('SIGTERM', () => gracefulShutdown('SIGTERM'));
// Handle Windows CTRL+C (SIGINT not always reliable on Windows)
if (process.platform === 'win32') {
const readline = require('readline');
const rl = readline.createInterface({
input: process.stdin,
output: process.stdout
});
rl.on('SIGINT', () => gracefulShutdown('SIGINT'));
}