/** * Chat WebSocket 路由处理器 * * 统一 WS 消息格式: * - connected: 连接确认 * - initialized: 初始化完成 (含 provider) * - thinking: 开始思考 * - acp_event: 所有内容事件(文本/工具/权限等) * - result: 轮次结束 * - error: 错误 * - disconnected: 断开 */ const { service } = require('../services/claude-chat-service'); const { extractUserFromToken, COOKIE_NAME } = require('../middleware/jwt-auth'); const config = require('../config'); const fs = require('fs'); const path = require('path'); /** * WebSocket 消息处理器 */ class ClaudeChatWSHandler { constructor(ws, wsId, user) { this.ws = ws; this.wsId = wsId; // 临时 WS 连接 ID this.session = null; // 关联的 Conversation (ChatSession) this.user = user; this.activeSkills = new Map(); this._skillsChanged = false; this._injectedSkillsHash = ''; this.ws.on('message', (data) => this.handleMessage(data)); this.ws.on('close', () => this.handleClose()); this.ws.on('error', (err) => this.handleError(err)); } /** * 发送消息给客户端 */ send(type, payload = {}) { if (this.ws.readyState === 1) { // OPEN this.ws.send(JSON.stringify({ type, ...payload })); } } /** * 处理收到的消息 */ async handleMessage(data) { try { const message = JSON.parse(data.toString()); const { type, ...payload } = message; switch (type) { case 'init': await this.handleInit(payload); break; case 'message': await this.handleSendMessage(payload); break; case 'activate_skill': this.handleActivateSkill(payload); break; case 'deactivate_skill': this.handleDeactivateSkill(payload); break; case 'disconnect': await this.handleDisconnect(); break; case 'getHistory': this.handleGetHistory(); break; case 'getProviders': this.handleGetProviders(); break; default: this.send('error', { error: 'Unknown message type' }); } } catch (err) { console.error('WebSocket message error:', err); this.send('error', { error: err.message }); } } /** * 初始化会话 — 接收 conversationId 复用已有对话 */ async handleInit({ workingDir = process.cwd(), provider = 'claude', conversationId } = {}) { try { // 如果客户端没传 conversationId,生成一个 const convId = conversationId || `conv-${Date.now().toString(36)}-${Math.random().toString(36).substr(2, 6)}`; // 获取或创建 Conversation this.session = service.getOrCreateConversation(convId, { userId: this.user.id }); // 重新绑定事件到当前 WS this.bindSessionEvents(); // 如果 Conversation 已有连接的 CLI,复用它 if (this.session.isConnected) { console.log(`[WS:${this.wsId}] 复用已有 CLI 连接 (conv: ${convId}, pid: ${this.session.pid})`); } else { // 需要新建 CLI 连接 await this.session.connect(workingDir, { provider }); // 如果有历史消息,注入上下文摘要让 Claude 记住之前的对话 if (this.session.messages.length > 0) { const contextSummary = this.session.getContextSummary(); if (contextSummary) { console.log(`[WS:${this.wsId}] 注入历史上下文 (${this.session.messages.length} msgs)`); const contextMsg = `以下是之前的对话历史摘要,请你记住这些上下文,然后等待用户的下一条消息:\n\n${contextSummary}\n\n---\n请回复"已了解,请继续"即可。`; await this.session.sendMessage(contextMsg); } } } this.send('initialized', { conversationId: convId, sessionId: convId, // 兼容旧前端 pid: this.session.pid, provider: this.session.provider, messageCount: this.session.messages.length }); } catch (err) { this.send('error', { error: err.message }); } } /** * 绑定 Conversation 事件到当前 WS * 每次新的 WS 连接都重新绑定,确保事件发到正确的 WS */ bindSessionEvents() { // 移除旧的事件监听器(避免重复绑定) this.session.removeAllListeners('acpEvent'); this.session.removeAllListeners('thinking'); this.session.removeAllListeners('messageComplete'); // 保留 error/disconnected/connected 的内部监听 // ── 唯一的内容通道 ── this.session.on('acpEvent', (event) => { this.send('acp_event', { event }); }); // ── 控制事件 ── this.session.on('thinking', () => { console.log(`[WS:${this.wsId}] thinking...`); this.send('thinking'); }); this.session.on('messageComplete', (result) => { console.log(`[WS:${this.wsId}] result:`, JSON.stringify(result)?.substring(0, 120)); this.send('result', { result }); }); } /** * 激活 Skill — 加载 SKILL.md + 关键参考文件 */ handleActivateSkill({ skillName }) { if (!skillName) { this.send('error', { error: 'skillName is required' }); return; } const skillDir = path.join(config.skillsDataDir, skillName); const skillMdPath = path.join(skillDir, 'SKILL.md'); if (!fs.existsSync(skillMdPath)) { this.send('error', { error: `Skill "${skillName}" not found` }); return; } // 加载完整 skill bundle(SKILL.md + 参考文件) const bundle = this._loadSkillBundle(skillName, skillDir); this.activeSkills.set(skillName, bundle); this._skillsChanged = true; console.log(`[WS:${this.wsId}] activated skill: ${skillName} (${bundle.length} chars)`); this.send('skill_activated', { skillName, activeSkills: [...this.activeSkills.keys()], }); } /** * 加载 Skill Bundle:SKILL.md + digest_summary.md + knowhow.md */ _loadSkillBundle(skillName, skillDir) { const parts = []; // 1. 主文件 SKILL.md(必须) const skillMd = fs.readFileSync(path.join(skillDir, 'SKILL.md'), 'utf-8'); parts.push(skillMd); // 2. 参考文件(可选,按优先级加载) const refFiles = [ { path: 'digest/digest_summary.md', label: 'Digest Summary' }, { path: 'hunter/knowhow.md', label: 'Domain Knowhow' }, { path: 'digest/benchmark.md', label: 'Quality Benchmark' }, ]; for (const ref of refFiles) { const refPath = path.join(skillDir, ref.path); if (fs.existsSync(refPath)) { try { const content = fs.readFileSync(refPath, 'utf-8'); // 限制单个参考文件不超过 8000 字符,防止 token 爆炸 const truncated = content.length > 8000 ? content.substring(0, 8000) + '\n\n... (truncated)' : content; parts.push(`\n### ${ref.label}\n${truncated}`); } catch { /* skip unreadable files */ } } } return parts.join('\n\n'); } /** * 取消激活 Skill */ handleDeactivateSkill({ skillName }) { this.activeSkills.delete(skillName); this._skillsChanged = true; console.log(`[WS:${this.wsId}] deactivated skill: ${skillName}`); this.send('skill_deactivated', { skillName, activeSkills: [...this.activeSkills.keys()], }); } /** * 构建 Skill 前缀(仅在 skills 变更时注入) * @returns {string} 需要注入的前缀,空字符串表示不需要注入 */ _buildSkillPrefix() { // 无 skill → 不注入 if (this.activeSkills.size === 0) return ''; // skills 没变化 → 不重复注入(AI 有上下文记忆) const currentHash = [...this.activeSkills.keys()].sort().join(','); if (!this._skillsChanged && this._injectedSkillsHash === currentHash) { return ''; } // 标记为已注入 this._injectedSkillsHash = currentHash; this._skillsChanged = false; const parts = []; for (const [name, content] of this.activeSkills) { parts.push(`\n${content}\n`); } return '请严格按照以下 Skill 指导来完成任务:\n\n' + parts.join('\n\n') + '\n\n---\n\n'; } /** * 构建附件上下文(将上传文件内容注入到消息中) */ _buildAttachmentContext(attachments) { if (!attachments || attachments.length === 0) return ''; const parts = []; for (const att of attachments) { const filePath = att.storedPath; if (!filePath || !fs.existsSync(filePath)) continue; try { if (att.fileType && att.fileType.startsWith('image/')) { // 图片 → base64 引用 const b64 = fs.readFileSync(filePath, 'base64'); const ext = path.extname(att.fileName).slice(1) || 'png'; parts.push( `\n` + `[Image: ${att.fileName} (${(att.fileSize / 1024).toFixed(1)}KB, ${ext})]\n` + `Base64: data:${att.fileType};base64,${b64}\n` + `` ); } else { // 文本/代码/CSV → 直接读取内容 let content = fs.readFileSync(filePath, 'utf-8'); // 限制单文件 20000 字符 if (content.length > 20000) { content = content.substring(0, 20000) + '\n\n... (truncated, total ' + content.length + ' chars)'; } parts.push( `\n` + `${content}\n` + `` ); } } catch (e) { console.warn(`[WS:${this.wsId}] Failed to read attachment: ${att.fileName}`, e.message); } } if (parts.length === 0) return ''; return '\n\n以下是用户上传的附件内容:\n\n' + parts.join('\n\n') + '\n\n---\n\n'; } /** * 发送消息(智能 Skill + 附件注入) */ async handleSendMessage({ text, skills, attachments }) { if (!this.session) { this.send('error', { error: 'Session not initialized' }); return; } // 如果消息附带了 skills 列表,先激活它们 if (skills && Array.isArray(skills)) { for (const sk of skills) { if (!this.activeSkills.has(sk)) { this.handleActivateSkill({ skillName: sk }); } } } // 构建完整消息: [skill 前缀] + 用户文本 + [附件上下文] const skillPrefix = this._buildSkillPrefix(); const attachContext = this._buildAttachmentContext(attachments); const fullText = (skillPrefix || '') + text + (attachContext || ''); try { await this.session.sendMessage(fullText); } catch (err) { this.send('error', { error: err.message }); } } /** * 断开连接 */ async handleDisconnect() { if (this.session) { this.session.disconnect(); } this.send('disconnected'); } /** * 获取消息历史 */ handleGetHistory() { if (!this.session) { this.send('error', { error: 'Session not initialized' }); return; } this.send('history', { messages: this.session.getMessages() }); } /** * 获取支持的 provider 列表 */ handleGetProviders() { this.send('providers', { providers: service.getProviders() }); } /** * 处理连接关闭 — 不销毁 Conversation,允许重连 */ handleClose() { console.log(`WebSocket ${this.wsId} closed (conv: ${this.session?.id || 'none'})`); // 不断开 Conversation,允许通过新 WS 重新连接 } /** * 处理错误 */ handleError(err) { console.error(`WebSocket error for ${this.wsId}:`, err); } } /** * 初始化 WebSocket 服务器 */ function initClaudeChatWSServer(wss) { wss.on('connection', (ws, req) => { // 从 WS handshake Cookie 中提取 JWT 鉴权 const cookieHeader = req.headers.cookie || ''; const cookies = {}; cookieHeader.split(';').forEach(c => { const [k, ...v] = c.trim().split('='); if (k) cookies[k] = v.join('='); }); const token = cookies[COOKIE_NAME]; const user = extractUserFromToken(token); if (!user) { console.warn('[WS] Unauthorized WebSocket connection rejected'); ws.close(4001, 'Unauthorized'); return; } // wsId 是临时的 WS 连接标识,不是 conversationId const wsId = 'ws-' + Date.now().toString(36) + Math.random().toString(36).substr(2, 5); console.log(`WebSocket connected: ${wsId} (user: ${user.id})`); const handler = new ClaudeChatWSHandler(ws, wsId, user); // 发送连接成功消息(客户端收到后发送 init + conversationId) handler.send('connected', { wsId }); }); console.log('Chat WebSocket server initialized'); } module.exports = { ClaudeChatWSHandler, initClaudeChatWSServer };