treesap
Version:
AI Agent Framework
332 lines • 12.5 kB
JavaScript
import { WebSocketServer, WebSocket } from 'ws';
import { v4 as uuidv4 } from 'uuid';
import { TerminalService } from './terminal.js';
export class WebSocketTerminalService {
static wss = null;
static clients = new Map();
static sessionClients = new Map(); // sessionId -> Set<clientId>
static initialize(server) {
this.wss = new WebSocketServer({
server,
path: '/terminal/ws'
});
this.wss.on('connection', (ws, request) => {
const clientId = uuidv4();
const client = {
id: clientId,
ws,
lastPing: new Date()
};
this.clients.set(clientId, client);
console.log(`WebSocket client connected: ${clientId}`);
// Set up message handler
ws.on('message', (data) => {
this.handleMessage(clientId, data);
});
// Set up disconnect handler
ws.on('close', () => {
this.handleDisconnect(clientId);
});
// Set up error handler
ws.on('error', (error) => {
console.error(`WebSocket error for client ${clientId}:`, error);
this.handleDisconnect(clientId);
});
// Send initial connection success
this.sendToClient(clientId, {
type: 'pong',
timestamp: Date.now()
});
// Set up ping/pong for connection health
this.setupPingPong(clientId);
});
console.log('WebSocket server initialized for terminal connections');
}
static setupPingPong(clientId) {
const client = this.clients.get(clientId);
if (!client)
return;
const pingInterval = setInterval(() => {
if (client.ws.readyState === WebSocket.OPEN) {
client.ws.ping();
client.lastPing = new Date();
}
else {
clearInterval(pingInterval);
}
}, 30000); // Ping every 30 seconds
client.ws.on('pong', () => {
client.lastPing = new Date();
});
client.ws.on('close', () => {
clearInterval(pingInterval);
});
}
static handleMessage(clientId, data) {
try {
const message = JSON.parse(data.toString());
const client = this.clients.get(clientId);
if (!client) {
console.error(`Client ${clientId} not found`);
return;
}
switch (message.type) {
case 'join':
this.handleJoin(clientId, message);
break;
case 'leave':
this.handleLeave(clientId, message);
break;
case 'input':
this.handleInput(clientId, message);
break;
case 'resize':
this.handleResize(clientId, message);
break;
case 'ping':
this.sendToClient(clientId, {
type: 'pong',
timestamp: Date.now()
});
break;
default:
console.warn(`Unknown message type: ${message.type}`);
}
}
catch (error) {
console.error(`Error parsing WebSocket message from ${clientId}:`, error);
}
}
static handleJoin(clientId, message) {
const client = this.clients.get(clientId);
if (!client || !message.sessionId)
return;
console.log(`Client ${clientId} joining session ${message.sessionId}`);
// Remove client from any existing session
if (client.sessionId) {
this.removeClientFromSession(clientId, client.sessionId);
}
// Get or create terminal session
let session = TerminalService.getSession(message.sessionId);
if (!session) {
console.log(`Creating new terminal session: ${message.sessionId}`);
session = TerminalService.createSession(message.sessionId);
}
// Update client info
client.sessionId = message.sessionId;
client.terminalId = message.terminalId;
// Add client to session tracking
this.addClientToSession(clientId, message.sessionId);
// Set up output listener for this session (if not already set up)
this.setupSessionOutputListener(message.sessionId);
// Send connection confirmation
this.sendToClient(clientId, {
type: 'connected',
sessionId: message.sessionId,
timestamp: Date.now()
});
// Notify all clients in this session about client count
this.broadcastClientCount(message.sessionId);
}
static handleLeave(clientId, message) {
const client = this.clients.get(clientId);
if (!client)
return;
console.log(`Client ${clientId} leaving session ${client.sessionId}`);
if (client.sessionId) {
this.removeClientFromSession(clientId, client.sessionId);
this.broadcastClientCount(client.sessionId);
}
client.sessionId = undefined;
client.terminalId = undefined;
}
static handleInput(clientId, message) {
const client = this.clients.get(clientId);
if (!client || !message.sessionId || message.data === undefined)
return;
;
// Get the terminal session
const session = TerminalService.getSession(message.sessionId);
if (!session) {
console.error(`Session ${message.sessionId} not found`);
this.sendToClient(clientId, {
type: 'error',
data: 'Terminal session not found',
timestamp: Date.now()
});
return;
}
// Send input directly to PTY (raw input like key presses)
try {
session.lastActivity = new Date();
session.process.write(message.data);
}
catch (error) {
console.error(`Failed to send input to session ${message.sessionId}:`, error);
this.sendToClient(clientId, {
type: 'error',
data: 'Failed to send input to terminal',
timestamp: Date.now()
});
}
}
static handleResize(clientId, message) {
const client = this.clients.get(clientId);
if (!client || !message.sessionId || message.cols === undefined || message.rows === undefined)
return;
// Get the terminal session
const session = TerminalService.getSession(message.sessionId);
if (!session) {
console.error(`Session ${message.sessionId} not found for resize`);
return;
}
// Resize the PTY
try {
console.log(`Resizing terminal session ${message.sessionId} to ${message.cols}x${message.rows}`);
session.process.resize(message.cols, message.rows);
session.lastActivity = new Date();
// Update session dimensions
session.cols = message.cols;
session.rows = message.rows;
}
catch (error) {
console.error(`Failed to resize session ${message.sessionId}:`, error);
}
}
static handleDisconnect(clientId) {
const client = this.clients.get(clientId);
console.log(`WebSocket client disconnected: ${clientId}`);
if (client?.sessionId) {
this.removeClientFromSession(clientId, client.sessionId);
this.broadcastClientCount(client.sessionId);
}
this.clients.delete(clientId);
}
static addClientToSession(clientId, sessionId) {
if (!this.sessionClients.has(sessionId)) {
this.sessionClients.set(sessionId, new Set());
}
this.sessionClients.get(sessionId).add(clientId);
}
static removeClientFromSession(clientId, sessionId) {
const clientSet = this.sessionClients.get(sessionId);
if (clientSet) {
clientSet.delete(clientId);
if (clientSet.size === 0) {
this.sessionClients.delete(sessionId);
// Clean up output listener if no clients are connected
this.cleanupSessionOutputListener(sessionId);
}
}
}
static setupSessionOutputListener(sessionId) {
const session = TerminalService.getSession(sessionId);
if (!session)
return;
// Check if listener already exists
if (session.eventEmitter.listenerCount('output') > 0) {
return; // Listener already set up
}
const handleOutput = (data) => {
this.broadcastToSession(sessionId, data);
};
session.eventEmitter.on('output', handleOutput);
console.log(`Set up output listener for session ${sessionId}`);
}
static cleanupSessionOutputListener(sessionId) {
const session = TerminalService.getSession(sessionId);
if (!session)
return;
session.eventEmitter.removeAllListeners('output');
console.log(`Cleaned up output listener for session ${sessionId}`);
}
static broadcastToSession(sessionId, data) {
const clientIds = this.sessionClients.get(sessionId);
if (!clientIds)
return;
const message = {
...data,
timestamp: Date.now()
};
for (const clientId of clientIds) {
this.sendToClient(clientId, message);
}
}
static broadcastClientCount(sessionId) {
const clientIds = this.sessionClients.get(sessionId);
const count = clientIds ? clientIds.size : 0;
if (clientIds) {
for (const clientId of clientIds) {
this.sendToClient(clientId, {
type: 'clients_count',
count,
timestamp: Date.now()
});
}
}
}
static sendToClient(clientId, message) {
const client = this.clients.get(clientId);
if (!client || client.ws.readyState !== WebSocket.OPEN) {
return;
}
try {
client.ws.send(JSON.stringify(message));
}
catch (error) {
console.error(`Error sending message to client ${clientId}:`, error);
this.handleDisconnect(clientId);
}
}
// API methods for external control
static getSessionClients(sessionId) {
const clientSet = this.sessionClients.get(sessionId);
return clientSet ? Array.from(clientSet) : [];
}
static sendCommandToSession(sessionId, command) {
// Send command to terminal
const success = TerminalService.executeCommand(sessionId, command);
if (success) {
// The output will be automatically broadcast to all connected clients
// via the output listener
console.log(`Command sent to session ${sessionId}: ${command.substring(0, 50)}...`);
}
return success;
}
static getActiveSessions() {
return Array.from(this.sessionClients.entries()).map(([sessionId, clients]) => ({
sessionId,
clientCount: clients.size
}));
}
static getConnectedClients() {
return this.clients.size;
}
static closeSession(sessionId) {
const clientIds = this.sessionClients.get(sessionId);
if (clientIds) {
// Notify all clients that the session is closing
for (const clientId of clientIds) {
this.sendToClient(clientId, {
type: 'session_closed',
sessionId,
timestamp: Date.now()
});
}
// Clean up client tracking
this.sessionClients.delete(sessionId);
}
// Clean up the terminal session
TerminalService.destroySession(sessionId);
}
static cleanup() {
if (this.wss) {
console.log('Closing WebSocket server...');
this.wss.close();
this.wss = null;
}
this.clients.clear();
this.sessionClients.clear();
}
}
//# sourceMappingURL=websocket.js.map