// AI生图队列管理器 const fs = require('fs'); const path = require('path'); const ReplaceCharacterHandler = require('./replace-character'); const { getDatabase } = require('./sql'); const QUEUE_FILE = path.join(__dirname, 'ai-queue.json'); // 超时配置(毫秒) const TASK_TIMEOUT = 5 * 60 * 1000; // 5分钟超时 // 队列状态 let queue = []; let isProcessing = false; let dbInstance = null; let timeoutCheckInterval = null; // 获取数据库实例 async function getDB() { if (!dbInstance) { dbInstance = await getDatabase(); } return dbInstance; } // 初始化:加载队列 function initQueue() { try { if (fs.existsSync(QUEUE_FILE)) { const data = fs.readFileSync(QUEUE_FILE, 'utf-8'); queue = JSON.parse(data); } } catch (error) { console.error('[AIQueue] 加载队列失败:', error); queue = []; } } // 保存队列 function saveQueue() { try { fs.writeFileSync(QUEUE_FILE, JSON.stringify(queue, null, 2), 'utf-8'); } catch (error) { console.error('[AIQueue] 保存队列失败:', error); } } // 添加任务到队列 async function addToQueue(username, taskData) { const taskId = `ai_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`; const task = { id: taskId, username: username.toLowerCase(), status: queue.length === 0 && !isProcessing ? 'rendering' : 'queued', createdAt: new Date().toISOString(), ...taskData }; queue.push(task); saveQueue(); // 添加到历史记录(不保存参考图) try { const db = await getDB(); db.addAIHistory(taskId, username, task.status, null); } catch (error) { console.error('[AIQueue] 添加历史记录失败:', error); } // 如果队列为空且没有正在处理的任务,立即开始处理 if (queue.length === 1 && !isProcessing) { processQueue(); } return taskId; } // 检查超时任务 async function checkTimeoutTasks() { try { const db = await getDB(); const timedOutTasks = db.checkAndMarkTimeoutTasks(TASK_TIMEOUT); if (timedOutTasks && timedOutTasks.length > 0) { console.log(`[AIQueue] 发现 ${timedOutTasks.length} 个超时任务,已标记为失败`); timedOutTasks.forEach(taskId => { console.log(`[AIQueue] 超时任务: ${taskId}`); }); } } catch (error) { console.error('[AIQueue] 检查超时任务失败:', error); } } // 启动超时检查定时器 function startTimeoutChecker() { if (timeoutCheckInterval) { clearInterval(timeoutCheckInterval); } // 每30秒检查一次超时任务 timeoutCheckInterval = setInterval(checkTimeoutTasks, 30000); console.log('[AIQueue] 超时检查定时器已启动(每30秒检查一次)'); } // 处理队列 async function processQueue() { if (isProcessing || queue.length === 0) { return; } isProcessing = true; while (queue.length > 0) { const task = queue[0]; // 更新状态为rendering,并记录开始时间 if (task.status === 'queued') { task.status = 'rendering'; task.renderStartTime = Date.now(); updateTaskStatus(task.id, 'rendering', null, null, null, task.renderStartTime); } try { console.log(`[AIQueue] 开始处理任务: ${task.id}`); // 调用Gemini API const result = await callGeminiAPIWithPromise( task.image1, task.image2, task.image1Width, task.image1Height, task.additionalPrompt || '' ); if (result.success && result.imageData) { // 保存图片到用户目录 const imageUrl = await saveAIImage(task.username, task.id, result.imageData); // 更新任务状态 task.status = 'completed'; task.imageUrl = imageUrl; task.completedAt = new Date().toISOString(); updateTaskStatus(task.id, 'completed', imageUrl); console.log(`[AIQueue] 任务完成: ${task.id}`); } else { // Gemini API 明确返回失败 const apiError = new Error(result.error || '生成失败'); apiError.isApiError = true; throw apiError; } } catch (error) { console.error(`[AIQueue] 任务失败: ${task.id}`, error); // 只有 Gemini API 明确返回失败时才标记为 failed if (error.isApiError) { task.status = 'failed'; task.error = error.message; task.completedAt = new Date().toISOString(); // 保存原始任务数据用于重试(不包含图片数据以节省空间,重试时从预览图重新加载) updateTaskStatus(task.id, 'failed', null, error.message, { image1Width: task.image1Width, image1Height: task.image1Height, additionalPrompt: task.additionalPrompt }); } else { // 其他错误(代码错误、网络错误等)自动重试 console.log(`[AIQueue] 非API错误,将任务重新加入队列末尾: ${task.id}`); task.status = 'queued'; task.retryCount = (task.retryCount || 0) + 1; // 最多重试3次 if (task.retryCount <= 3) { queue.push({ ...task }); updateTaskStatus(task.id, 'queued'); } else { console.error(`[AIQueue] 任务重试次数超过限制,标记为失败: ${task.id}`); task.status = 'failed'; task.error = '多次重试失败:' + error.message; task.completedAt = new Date().toISOString(); updateTaskStatus(task.id, 'failed', null, task.error, { image1Width: task.image1Width, image1Height: task.image1Height, additionalPrompt: task.additionalPrompt }); } } } // 从队列中移除 queue.shift(); saveQueue(); } isProcessing = false; } // 调用Gemini API(Promise版本) function callGeminiAPIWithPromise(image1Base64, image2Base64, image1Width, image1Height, additionalPrompt) { return new Promise((resolve, reject) => { const https = require('https'); const startTime = Date.now(); // 移除 data:image/png;base64, 前缀(如果有) const cleanImage1 = image1Base64.replace(/^data:image\/\w+;base64,/, ''); const cleanImage2 = image2Base64.replace(/^data:image\/\w+;base64,/, ''); // 构建请求内容 const promptText = ReplaceCharacterHandler.buildPromptText(image1Width, image1Height, additionalPrompt); const content = [ { type: "image_url", image_url: { url: `data:image/png;base64,${cleanImage1}` } }, { type: "image_url", image_url: { url: `data:image/png;base64,${cleanImage2}` } }, { type: "text", text: promptText } ]; const requestData = JSON.stringify({ model: "gemini-3-pro-image-preview", messages: [{ role: "user", content: content }] }); const options = { hostname: 'api.chatanywhere.tech', port: 443, path: '/v1/chat/completions', method: 'POST', headers: { 'Authorization': 'Bearer sk-j32LgDixK6pfESYGfJtgc2Tzlmszx5NZhSH0sOzpLQkYuKek', 'Content-Type': 'application/json', 'Content-Length': Buffer.byteLength(requestData) }, timeout: 300000 // 5分钟超时 }; // 请求日志 console.log('\n========== [Gemini API 请求] =========='); console.log(`[时间] ${new Date().toLocaleString()}`); console.log(`[模型] gemini-3-pro-image-preview`); console.log(`[目标] https://${options.hostname}${options.path}`); console.log(`[图片1大小] ${(cleanImage1.length / 1024).toFixed(2)} KB`); console.log(`[图片2大小] ${(cleanImage2.length / 1024).toFixed(2)} KB`); console.log(`[输出尺寸] ${image1Width} x ${image1Height}`); console.log(`[附加提示] ${additionalPrompt || '(无)'}`); console.log(`[请求体大小] ${(Buffer.byteLength(requestData) / 1024).toFixed(2)} KB`); console.log('========================================\n'); const geminiReq = https.request(options, (geminiRes) => { let responseData = ''; geminiRes.on('data', (chunk) => { responseData += chunk; }); geminiRes.on('end', () => { const elapsed = ((Date.now() - startTime) / 1000).toFixed(2); console.log('\n========== [Gemini API 响应] =========='); console.log(`[时间] ${new Date().toLocaleString()}`); console.log(`[耗时] ${elapsed} 秒`); console.log(`[状态码] ${geminiRes.statusCode}`); console.log(`[响应大小] ${(responseData.length / 1024).toFixed(2)} KB`); try { if (geminiRes.statusCode !== 200) { console.log(`[错误] API返回非200状态`); console.log(`[响应内容] ${responseData.substring(0, 500)}...`); console.log('========================================\n'); reject(new Error(`Gemini API error: ${geminiRes.statusCode}`)); return; } const response = JSON.parse(responseData); // 打印响应结构 console.log(`[响应结构] choices: ${response.choices?.length || 0}`); if (response.choices && response.choices[0]) { const msg = response.choices[0].message; if (msg && msg.content) { if (Array.isArray(msg.content)) { console.log(`[内容类型] 数组,包含 ${msg.content.length} 个元素`); msg.content.forEach((item, i) => { if (item.type === 'text') { console.log(` [${i}] text: ${item.text?.substring(0, 100) || '(空)'}...`); } else if (item.type === 'image_url') { const imgData = item.image_url?.url || ''; console.log(` [${i}] image: ${(imgData.length / 1024).toFixed(2)} KB`); } else { console.log(` [${i}] ${item.type}: ...`); } }); } else { console.log(`[内容类型] 字符串,长度 ${msg.content.length}`); } } } if (response.usage) { console.log(`[Token使用] prompt: ${response.usage.prompt_tokens}, completion: ${response.usage.completion_tokens}, total: ${response.usage.total_tokens}`); } // 解析响应,提取图片 const imageData = ReplaceCharacterHandler.extractImageFromResponse(response); if (!imageData) { console.log(`[结果] ✗ 无法从响应中提取图片`); console.log('========================================\n'); reject(new Error('Failed to extract image from response')); return; } console.log(`[结果] ✓ 成功提取图片,大小: ${(imageData.length / 1024).toFixed(2)} KB`); console.log('========================================\n'); resolve({ success: true, imageData: imageData }); } catch (error) { console.log(`[错误] 解析响应失败: ${error.message}`); console.log('========================================\n'); reject(error); } }); }); geminiReq.on('error', (error) => { const elapsed = ((Date.now() - startTime) / 1000).toFixed(2); console.log('\n========== [Gemini API 错误] =========='); console.log(`[时间] ${new Date().toLocaleString()}`); console.log(`[耗时] ${elapsed} 秒`); console.log(`[错误类型] ${error.code || 'Unknown'}`); console.log(`[错误信息] ${error.message}`); console.log('========================================\n'); reject(error); }); geminiReq.on('timeout', () => { const elapsed = ((Date.now() - startTime) / 1000).toFixed(2); console.log('\n========== [Gemini API 超时] =========='); console.log(`[时间] ${new Date().toLocaleString()}`); console.log(`[耗时] ${elapsed} 秒`); console.log(`[超时设置] ${options.timeout / 1000} 秒`); console.log('========================================\n'); geminiReq.destroy(); reject(new Error('Request timeout')); }); geminiReq.write(requestData); geminiReq.end(); }); } // 保存AI生成的图片 function saveAIImage(username, taskId, imageBase64) { const usersDir = path.join(__dirname, 'users'); const userDir = path.join(usersDir, username.toLowerCase()); const aiDir = path.join(userDir, 'ai-images'); // 确保目录存在 if (!fs.existsSync(aiDir)) { fs.mkdirSync(aiDir, { recursive: true }); } const imagePath = path.join(aiDir, `${taskId}.png`); const imageBuffer = Buffer.from(imageBase64, 'base64'); fs.writeFileSync(imagePath, imageBuffer); return `/api/ai/image?username=${encodeURIComponent(username)}&id=${encodeURIComponent(taskId)}`; } // 更新任务状态 async function updateTaskStatus(taskId, status, imageUrl = null, error = null, retryData = null, renderStartTime = null) { try { const db = await getDB(); db.updateAITaskStatus(taskId, status, imageUrl, error, renderStartTime); if (retryData) { db.updateAITaskRetryData(taskId, retryData); } } catch (err) { console.error('[AIQueue] 更新任务状态失败:', err); } } // 获取用户AI历史 async function getUserAIHistory(username) { try { const db = await getDB(); return db.getAIHistory(username); } catch (error) { console.error('[AIQueue] 获取用户AI历史失败:', error); return []; } } // 处理AI生图请求(队列版本) function handleAIRequest(req, res) { if (req.method !== 'POST') { res.writeHead(405, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: 'Method not allowed' })); return; } let body = ''; req.on('data', (chunk) => { body += chunk.toString(); }); req.on('end', async () => { try { const data = JSON.parse(body); const { username, image1, image2, image1Width, image1Height, additionalPrompt } = data; if (!username) { res.writeHead(400, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ success: false, error: '缺少用户名参数' })); return; } if (!image1 || !image2) { res.writeHead(400, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ success: false, error: 'Missing required fields: image1, image2' })); return; } if (!image1Width || !image1Height) { res.writeHead(400, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ success: false, error: 'Missing required fields: image1Width, image1Height' })); return; } // 添加到队列 const taskId = await addToQueue(username, { image1, image2, image1Width, image1Height, additionalPrompt: additionalPrompt || '' }); // 立即返回任务ID res.writeHead(200, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ success: true, taskId: taskId, message: '请求生图成功,正在处理中...' })); // 异步处理队列 processQueue(); } catch (error) { console.error('[AIQueue] 处理请求失败:', error); res.writeHead(500, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ success: false, error: '处理失败', details: error.message })); } }); req.on('error', (error) => { console.error('[AIQueue] 请求错误:', error); res.writeHead(500, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ success: false, error: 'Request error', details: error.message })); }); } // 重试失败的任务 - 直接删除失败记录(不再支持重试,因为不保存参考图) async function retryTask(taskId, username) { try { const db = await getDB(); const task = db.getAITask(taskId); if (!task) { return { success: false, error: '任务不存在' }; } if (task.status !== 'failed') { return { success: false, error: '只能删除失败的任务' }; } // 直接删除失败的任务记录 db.deleteAITask(taskId); return { success: true, message: '已删除失败记录' }; } catch (error) { console.error('[AIQueue] 删除任务失败:', error); return { success: false, error: error.message }; } } // 处理重试请求 function handleRetryRequest(req, res) { let body = ''; req.on('data', chunk => { body += chunk.toString(); }); req.on('end', async () => { try { const { taskId, username } = JSON.parse(body); if (!taskId || !username) { res.writeHead(400, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ success: false, error: '缺少必要参数' })); return; } const result = await retryTask(taskId, username); res.writeHead(result.success ? 200 : 400, { 'Content-Type': 'application/json' }); res.end(JSON.stringify(result)); } catch (error) { console.error('[AIQueue] 重试请求失败:', error); res.writeHead(500, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ success: false, error: '处理失败', details: error.message })); } }); } // 初始化 initQueue(); startTimeoutChecker(); module.exports = { handleAIRequest, handleRetryRequest, getUserAIHistory, processQueue, TASK_TIMEOUT };