/**
* 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
};