132 lines
3.3 KiB
JavaScript
132 lines
3.3 KiB
JavaScript
import WebSocket, { WebSocketServer } from 'ws';
|
|
import http from 'http';
|
|
import { DBModel } from "./models/index.js";
|
|
import { chatTask } from "./agent/task.js";
|
|
|
|
export default class WebSocketServerManager {
|
|
constructor(port = 8080) {
|
|
this.port = port;
|
|
this.wss = null;
|
|
this.server = null;
|
|
}
|
|
|
|
start() {
|
|
this.server = http.createServer((req, res) => {
|
|
res.writeHead(426, { 'Content-Type': 'text/plain' });
|
|
res.end('Upgrade Required: WebSocket expected');
|
|
});
|
|
|
|
this.wss = new WebSocketServer({
|
|
server: this.server,
|
|
perMessageDeflate: {
|
|
zlibDeflateOptions: {
|
|
chunkSize: 1024,
|
|
memLevel: 7,
|
|
level: 3
|
|
},
|
|
zlibInflateOptions: {
|
|
chunkSize: 10 * 1024
|
|
},
|
|
clientNoContextTakeover: true,
|
|
serverNoContextTakeover: true,
|
|
serverMaxWindowBits: 10,
|
|
concurrencyLimit: 10,
|
|
threshold: 1024
|
|
}
|
|
});
|
|
|
|
this.wss.on('connection', (ws, req) => {
|
|
const path = req.url;
|
|
console.log(`WebSocket connected: ${req.socket.remoteAddress}, path: ${path}`);
|
|
|
|
if (path === '/chat') {
|
|
this._handleChatConnection(ws);
|
|
} else {
|
|
ws.send(JSON.stringify({ type: 'system', content: 'Unknown path' }));
|
|
ws.close(1000, 'Unknown path');
|
|
}
|
|
});
|
|
|
|
this.server.listen(this.port, () => {
|
|
console.log(`WebSocket server started on ws://localhost:${this.port}`);
|
|
});
|
|
}
|
|
|
|
_handleChatConnection(ws) {
|
|
ws.send(JSON.stringify({ type: 'system', content: 'Connected to chat server' }));
|
|
|
|
ws.on('message', (data) => this._handleMessage(ws, data));
|
|
|
|
ws.on('close', () => {
|
|
console.log('Chat client disconnected');
|
|
});
|
|
|
|
ws.on('error', (err) => {
|
|
console.error('WebSocket error:', err);
|
|
});
|
|
}
|
|
|
|
async _handleMessage(ws, data) {
|
|
try {
|
|
const raw = data.toString();
|
|
console.log('Received:', raw);
|
|
|
|
const msg = JSON.parse(raw);
|
|
|
|
if (msg.type === 'ping') {
|
|
ws.send(JSON.stringify({ type: 'pong' }));
|
|
return;
|
|
}
|
|
|
|
if (msg.type === 'chat' || msg.type === 'clear') {
|
|
const userInfo = await this.getUserInfo(msg.userId);
|
|
chatTask.streamChat(userInfo, msg, (source, type, content, id) => {
|
|
ws.send(JSON.stringify({ source, type, content, id }));
|
|
});
|
|
return;
|
|
}
|
|
|
|
if (msg.type === 'system') {
|
|
console.log('[System]', msg.content);
|
|
return;
|
|
}
|
|
|
|
console.log('[Unknown type]', msg);
|
|
} catch (e) {
|
|
console.error('Message parse error:', e);
|
|
ws.send(JSON.stringify({ type: 'error', content: 'Invalid JSON' }));
|
|
}
|
|
}
|
|
|
|
broadcast(message) {
|
|
if (!this.wss) return;
|
|
this.wss.clients.forEach((client) => {
|
|
if (client.readyState === WebSocket.OPEN) {
|
|
client.send(message);
|
|
}
|
|
});
|
|
}
|
|
|
|
stop() {
|
|
if (this.wss) {
|
|
this.wss.close();
|
|
this.wss = null;
|
|
}
|
|
if (this.server) {
|
|
this.server.close();
|
|
this.server = null;
|
|
}
|
|
console.log('WebSocket server stopped');
|
|
}
|
|
|
|
getClientsCount() {
|
|
if (!this.wss) return 0;
|
|
return this.wss.clients.size;
|
|
}
|
|
|
|
async getUserInfo(userId) {
|
|
if (!DBModel.User || userId.length === 0) return null;
|
|
return await DBModel.User.findById(userId).lean().exec();
|
|
}
|
|
}
|