jb/paseo
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: encapsulate WebSocket client state into ClientSession class

Browse Source
This commit is contained in:
Mohamed Boudra
2025-10-20 14:29:47 +02:00
parent 30fc716b02
commit 8365fb889b
1 changed files with 148 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/voice-assistant/src/server/websocket-server.ts
+148 -181
View File
@@ -11,23 +11,68 @@ import { createConversation, deleteConversation } from "./agent/orchestrator.js"
type ProcessingPhase = 'idle' | 'transcribing' | 'llm'; type ProcessingPhase = 'idle' | 'transcribing' | 'llm';
/**
* Encapsulates all state for a single WebSocket client session
*/
class ClientSession {
public readonly clientId: string;
public readonly conversationId: string;
public abortController: AbortController;
public processingPhase: ProcessingPhase = 'idle';
public pendingAudioSegments: Array<{audio: Buffer, format: string}> = [];
public bufferTimeout: NodeJS.Timeout | null = null;
public audioBuffer: { chunks: Buffer[]; format: string } | null = null;
constructor(clientId: string, conversationId: string) {
this.clientId = clientId;
this.conversationId = conversationId;
this.abortController = new AbortController();
}
/**
* Create a new AbortController, aborting the previous one
*/
public createAbortController(): AbortController {
this.abortController.abort();
this.abortController = new AbortController();
return this.abortController;
}
/**
* Set the processing phase for this session
*/
public setPhase(phase: ProcessingPhase): void {
this.processingPhase = phase;
console.log(`[Session ${this.clientId}] Phase set to '${phase}'`);
}
/**
* Clean up all session resources
*/
public cleanup(): void {
// Abort any ongoing operations
this.abortController.abort();
// Clear buffer timeout
if (this.bufferTimeout) {
clearTimeout(this.bufferTimeout);
this.bufferTimeout = null;
}
// Clear buffers
this.pendingAudioSegments = [];
this.audioBuffer = null;
}
}
export class VoiceAssistantWebSocketServer { export class VoiceAssistantWebSocketServer {
private wss: WebSocketServer; private wss: WebSocketServer;
private clients: Map<WebSocket, string> = new Map(); // Map ws to client ID private sessions: Map<WebSocket, ClientSession> = new Map();
private conversationIds: Map<WebSocket, string> = new Map(); // Map ws to conversation ID private conversationIdToWs: Map<string, WebSocket> = new Map(); // Reverse map for phase management
private conversationIdToWs: Map<string, WebSocket> = new Map(); // Reverse map: conversation ID to ws
private abortControllers: Map<WebSocket, AbortController> = new Map(); // Map ws to AbortController
private messageHandler?: (conversationId: string, message: string, abortSignal: AbortSignal) => Promise<void>; private messageHandler?: (conversationId: string, message: string, abortSignal: AbortSignal) => Promise<void>;
private audioHandler?: (conversationId: string, audio: Buffer, format: string, abortSignal: AbortSignal) => Promise<string>; private audioHandler?: (conversationId: string, audio: Buffer, format: string, abortSignal: AbortSignal) => Promise<string>;
private audioBuffers: Map<string, { chunks: Buffer[]; format: string }> =
new Map();
private clientIdCounter: number = 0; private clientIdCounter: number = 0;
// Audio buffering for interruption handling
private processingPhases: Map<WebSocket, ProcessingPhase> = new Map();
private pendingAudioSegments: Map<WebSocket, Array<{audio: Buffer, format: string}>> = new Map();
private bufferTimeouts: Map<WebSocket, NodeJS.Timeout> = new Map();
constructor(server: HTTPServer) { constructor(server: HTTPServer) {
this.wss = new WebSocketServer({ server, path: "/ws" }); this.wss = new WebSocketServer({ server, path: "/ws" });
@@ -39,25 +84,17 @@ export class VoiceAssistantWebSocketServer {
} }
private handleConnection(ws: WebSocket): void { private handleConnection(ws: WebSocket): void {
// Generate unique client ID // Generate unique client ID and conversation
const clientId = `client-${++this.clientIdCounter}`; const clientId = `client-${++this.clientIdCounter}`;
this.clients.set(ws, clientId);
// Create new conversation for this client
const conversationId = createConversation(); const conversationId = createConversation();
this.conversationIds.set(ws, conversationId);
this.conversationIdToWs.set(conversationId, ws); // Add reverse mapping
// Create AbortController for this client // Create session for this client
const abortController = new AbortController(); const session = new ClientSession(clientId, conversationId);
this.abortControllers.set(ws, abortController); this.sessions.set(ws, session);
this.conversationIdToWs.set(conversationId, ws);
// Initialize processing state
this.processingPhases.set(ws, 'idle');
this.pendingAudioSegments.set(ws, []);
console.log( console.log(
`[WS] Client connected: ${clientId} with conversation ${conversationId} (total: ${this.clients.size})` `[WS] Client connected: ${clientId} with conversation ${conversationId} (total: ${this.sessions.size})`
); );
// Send welcome message // Send welcome message
@@ -74,75 +111,34 @@ export class VoiceAssistantWebSocketServer {
}); });
ws.on("close", () => { ws.on("close", () => {
const clientId = this.clients.get(ws); const session = this.sessions.get(ws);
const conversationId = this.conversationIds.get(ws); if (!session) return;
const abortController = this.abortControllers.get(ws);
// Abort any ongoing operations console.log(`[WS] Client disconnected: ${session.clientId} (total: ${this.sessions.size - 1})`);
if (abortController) {
abortController.abort();
this.abortControllers.delete(ws);
console.log(`[WS] Aborted operations for ${clientId}`);
}
// Clear buffer timeout // Clean up session
const timeout = this.bufferTimeouts.get(ws); session.cleanup();
if (timeout) { deleteConversation(session.conversationId);
clearTimeout(timeout);
this.bufferTimeouts.delete(ws);
}
// Clean up state tracking // Remove from maps
this.processingPhases.delete(ws); this.sessions.delete(ws);
this.pendingAudioSegments.delete(ws); this.conversationIdToWs.delete(session.conversationId);
if (clientId) { console.log(`[WS] Conversation ${session.conversationId} deleted`);
this.clients.delete(ws);
console.log(
`[WS] Client disconnected: ${clientId} (total: ${this.clients.size})`
);
}
if (conversationId) {
deleteConversation(conversationId);
this.conversationIds.delete(ws);
this.conversationIdToWs.delete(conversationId); // Clean up reverse mapping
console.log(`[WS] Conversation ${conversationId} deleted`);
}
}); });
ws.on("error", (error) => { ws.on("error", (error) => {
console.error("[WS] Client error:", error); console.error("[WS] Client error:", error);
const clientId = this.clients.get(ws); const session = this.sessions.get(ws);
const conversationId = this.conversationIds.get(ws); if (!session) return;
const abortController = this.abortControllers.get(ws);
// Abort any ongoing operations // Clean up session
if (abortController) { session.cleanup();
abortController.abort(); deleteConversation(session.conversationId);
this.abortControllers.delete(ws);
}
// Clear buffer timeout // Remove from maps
const timeout = this.bufferTimeouts.get(ws); this.sessions.delete(ws);
if (timeout) { this.conversationIdToWs.delete(session.conversationId);
clearTimeout(timeout);
this.bufferTimeouts.delete(ws);
}
// Clean up state tracking
this.processingPhases.delete(ws);
this.pendingAudioSegments.delete(ws);
if (clientId) {
this.clients.delete(ws);
}
if (conversationId) {
deleteConversation(conversationId);
this.conversationIds.delete(ws);
this.conversationIdToWs.delete(conversationId); // Clean up reverse mapping
}
}); });
} }
@@ -163,24 +159,17 @@ export class VoiceAssistantWebSocketServer {
case "user_message": case "user_message":
// Handle user message through orchestrator // Handle user message through orchestrator
const payload = message.payload as { message: string }; const payload = message.payload as { message: string };
const conversationId = this.conversationIds.get(ws); const session = this.sessions.get(ws);
if (!conversationId) { if (!session) {
console.error("[WS] No conversation found for client"); console.error("[WS] No session found for client");
break; break;
} }
// Abort any ongoing request // Create new abort controller for this request (aborts previous one)
const oldController = this.abortControllers.get(ws); const newController = session.createAbortController();
if (oldController) {
oldController.abort();
}
// Create new abort controller for this request
const newController = new AbortController();
this.abortControllers.set(ws, newController);
if (this.messageHandler) { if (this.messageHandler) {
await this.messageHandler(conversationId, payload.message, newController.signal); await this.messageHandler(session.conversationId, payload.message, newController.signal);
} else { } else {
console.warn("[WS] No message handler registered"); console.warn("[WS] No message handler registered");
} }
@@ -211,31 +200,26 @@ export class VoiceAssistantWebSocketServer {
} }
private async handleAbortRequest(ws: WebSocket): Promise<void> { private async handleAbortRequest(ws: WebSocket): Promise<void> {
const phase = this.processingPhases.get(ws); const session = this.sessions.get(ws);
const abortController = this.abortControllers.get(ws); if (!session) return;
console.log(`[WS] Abort request received, current phase: ${phase}`); console.log(`[WS] Abort request received, current phase: ${session.processingPhase}`);
if (phase === 'llm') { if (session.processingPhase === 'llm') {
// Already in LLM phase - abort immediately // Already in LLM phase - abort immediately
if (abortController) { session.abortController.abort();
abortController.abort(); console.log(`[WS] Aborted LLM processing`);
console.log(`[WS] Aborted LLM processing`);
}
// Reset phase to idle // Reset phase to idle
this.processingPhases.set(ws, 'idle'); session.setPhase('idle');
// Clear any pending segments from previous interruptions // Clear any pending segments and timeouts
this.pendingAudioSegments.set(ws, []); session.pendingAudioSegments = [];
if (session.bufferTimeout) {
// Clear any pending timeout clearTimeout(session.bufferTimeout);
const timeout = this.bufferTimeouts.get(ws); session.bufferTimeout = null;
if (timeout) {
clearTimeout(timeout);
this.bufferTimeouts.delete(ws);
} }
} else if (phase === 'transcribing') { } else if (session.processingPhase === 'transcribing') {
// Still in STT phase - we'll buffer the next audio // Still in STT phase - we'll buffer the next audio
// Don't abort yet, just set a flag by keeping the current abort controller // Don't abort yet, just set a flag by keeping the current abort controller
console.log(`[WS] Will buffer next audio segment (currently transcribing)`); console.log(`[WS] Will buffer next audio segment (currently transcribing)`);
@@ -253,7 +237,7 @@ export class VoiceAssistantWebSocketServer {
public broadcast(message: WebSocketMessage): void { public broadcast(message: WebSocketMessage): void {
const payload = JSON.stringify(message); const payload = JSON.stringify(message);
this.clients.forEach((_clientId, client) => { this.sessions.forEach((_session, client) => {
if (client.readyState === 1) { if (client.readyState === 1) {
// WebSocket.OPEN = 1 // WebSocket.OPEN = 1
client.send(payload); client.send(payload);
@@ -291,8 +275,10 @@ export class VoiceAssistantWebSocketServer {
public setPhaseForConversation(conversationId: string, phase: ProcessingPhase): void { public setPhaseForConversation(conversationId: string, phase: ProcessingPhase): void {
const ws = this.conversationIdToWs.get(conversationId); const ws = this.conversationIdToWs.get(conversationId);
if (ws) { if (ws) {
this.processingPhases.set(ws, phase); const session = this.sessions.get(ws);
console.log(`[WS] Phase set to '${phase}' for conversation ${conversationId}`); if (session) {
session.setPhase(phase);
}
} else { } else {
console.warn(`[WS] Cannot set phase for unknown conversation ${conversationId}`); console.warn(`[WS] Cannot set phase for unknown conversation ${conversationId}`);
} }
@@ -303,17 +289,9 @@ export class VoiceAssistantWebSocketServer {
payload: AudioChunkPayload payload: AudioChunkPayload
): Promise<void> { ): Promise<void> {
try { try {
// Use client-specific key for buffering const session = this.sessions.get(ws);
const clientId = this.clients.get(ws); if (!session) {
if (!clientId) { console.error("[WS] No session found for WebSocket");
console.error("[WS] No client ID found for WebSocket");
return;
}
// Get conversation ID for this client
const conversationId = this.conversationIds.get(ws);
if (!conversationId) {
console.error("[WS] No conversation found for client");
return; return;
} }
@@ -322,24 +300,22 @@ export class VoiceAssistantWebSocketServer {
if (!payload.isLast) { if (!payload.isLast) {
// Buffer the chunk // Buffer the chunk
if (!this.audioBuffers.has(clientId)) { if (!session.audioBuffer) {
this.audioBuffers.set(clientId, { session.audioBuffer = {
chunks: [], chunks: [],
format: payload.format, format: payload.format,
}); };
} }
const buffer = this.audioBuffers.get(clientId)!; session.audioBuffer.chunks.push(audioBuffer);
buffer.chunks.push(audioBuffer);
console.log( console.log(
`[WS] Buffered audio chunk (${audioBuffer.length} bytes, total chunks: ${buffer.chunks.length})` `[WS] Buffered audio chunk (${audioBuffer.length} bytes, total chunks: ${session.audioBuffer.chunks.length})`
); );
} else { } else {
// Last chunk - complete audio segment received // Last chunk - complete audio segment received
const buffer = this.audioBuffers.get(clientId); const allChunks = session.audioBuffer
const allChunks = buffer ? [...session.audioBuffer.chunks, audioBuffer]
? [...buffer.chunks, audioBuffer]
: [audioBuffer]; : [audioBuffer];
const format = buffer?.format || payload.format; const format = session.audioBuffer?.format || payload.format;
// Concatenate all chunks for this segment // Concatenate all chunks for this segment
const currentSegmentAudio = Buffer.concat(allChunks); const currentSegmentAudio = Buffer.concat(allChunks);
@@ -348,34 +324,29 @@ export class VoiceAssistantWebSocketServer {
); );
// Clear chunk buffer // Clear chunk buffer
this.audioBuffers.delete(clientId); session.audioBuffer = null;
// Get current phase and pending segments
const currentPhase = this.processingPhases.get(ws) || 'idle';
const pendingSegments = this.pendingAudioSegments.get(ws) || [];
// Decision: buffer or process? // Decision: buffer or process?
const shouldBuffer = currentPhase === 'transcribing' && pendingSegments.length === 0; const shouldBuffer = session.processingPhase === 'transcribing' && session.pendingAudioSegments.length === 0;
if (shouldBuffer) { if (shouldBuffer) {
// Currently transcribing first segment - buffer this one // Currently transcribing first segment - buffer this one
console.log(`[WS] Buffering audio segment (phase: ${currentPhase})`); console.log(`[WS] Buffering audio segment (phase: ${session.processingPhase})`);
pendingSegments.push({ audio: currentSegmentAudio, format }); session.pendingAudioSegments.push({ audio: currentSegmentAudio, format });
this.pendingAudioSegments.set(ws, pendingSegments);
// Set timeout to process buffer if no more audio arrives // Set timeout to process buffer if no more audio arrives
this.setBufferTimeout(ws, conversationId); this.setBufferTimeout(ws, session.conversationId);
} else if (pendingSegments.length > 0) { } else if (session.pendingAudioSegments.length > 0) {
// We have buffered segments - add this one and process all together // We have buffered segments - add this one and process all together
pendingSegments.push({ audio: currentSegmentAudio, format }); session.pendingAudioSegments.push({ audio: currentSegmentAudio, format });
console.log(`[WS] Processing ${pendingSegments.length} buffered audio segments together`); console.log(`[WS] Processing ${session.pendingAudioSegments.length} buffered audio segments together`);
// Clear pending segments and timeout // Clear pending segments and timeout
this.pendingAudioSegments.set(ws, []); const pendingSegments = [...session.pendingAudioSegments];
const timeout = this.bufferTimeouts.get(ws); session.pendingAudioSegments = [];
if (timeout) { if (session.bufferTimeout) {
clearTimeout(timeout); clearTimeout(session.bufferTimeout);
this.bufferTimeouts.delete(ws); session.bufferTimeout = null;
} }
// Concatenate all segments // Concatenate all segments
@@ -383,10 +354,10 @@ export class VoiceAssistantWebSocketServer {
const combinedAudio = Buffer.concat(allSegmentAudios); const combinedAudio = Buffer.concat(allSegmentAudios);
// Process combined audio // Process combined audio
await this.processAudio(ws, conversationId, combinedAudio, format); await this.processAudio(ws, session.conversationId, combinedAudio, format);
} else { } else {
// Normal flow - no buffering needed // Normal flow - no buffering needed
await this.processAudio(ws, conversationId, currentSegmentAudio, format); await this.processAudio(ws, session.conversationId, currentSegmentAudio, format);
} }
} }
} catch (error: any) { } catch (error: any) {
@@ -406,18 +377,14 @@ export class VoiceAssistantWebSocketServer {
audio: Buffer, audio: Buffer,
format: string format: string
): Promise<void> { ): Promise<void> {
// Abort any ongoing request const session = this.sessions.get(ws);
const oldAbortController = this.abortControllers.get(ws); if (!session) return;
if (oldAbortController) {
oldAbortController.abort();
}
// Create new abort controller for this request // Create new abort controller for this request (aborts previous one)
const newAbortController = new AbortController(); const newAbortController = session.createAbortController();
this.abortControllers.set(ws, newAbortController);
// Set phase to transcribing // Set phase to transcribing
this.processingPhases.set(ws, 'transcribing'); session.setPhase('transcribing');
if (this.audioHandler) { if (this.audioHandler) {
this.broadcastActivityLog({ this.broadcastActivityLog({
@@ -440,47 +407,47 @@ export class VoiceAssistantWebSocketServer {
// It will set phase to 'llm' before LLM processing and 'idle' after completion // It will set phase to 'llm' before LLM processing and 'idle' after completion
} catch (error: any) { } catch (error: any) {
// If error occurs, reset to idle as safety net // If error occurs, reset to idle as safety net
this.processingPhases.set(ws, 'idle'); session.setPhase('idle');
throw error; throw error;
} }
} else { } else {
console.warn("[WS] No audio handler registered"); console.warn("[WS] No audio handler registered");
this.processingPhases.set(ws, 'idle'); session.setPhase('idle');
} }
} }
private setBufferTimeout(ws: WebSocket, conversationId: string): void { private setBufferTimeout(ws: WebSocket, conversationId: string): void {
const session = this.sessions.get(ws);
if (!session) return;
// Clear any existing timeout // Clear any existing timeout
const existingTimeout = this.bufferTimeouts.get(ws); if (session.bufferTimeout) {
if (existingTimeout) { clearTimeout(session.bufferTimeout);
clearTimeout(existingTimeout);
} }
// Set new timeout (10 seconds) // Set new timeout (10 seconds)
const timeout = setTimeout(async () => { session.bufferTimeout = setTimeout(async () => {
console.log(`[WS] Buffer timeout reached, processing pending segments`); console.log(`[WS] Buffer timeout reached, processing pending segments`);
const pendingSegments = this.pendingAudioSegments.get(ws) || []; if (session.pendingAudioSegments.length > 0) {
if (pendingSegments.length > 0) {
// Concatenate all pending segments // Concatenate all pending segments
const pendingSegments = [...session.pendingAudioSegments];
const allSegmentAudios = pendingSegments.map(s => s.audio); const allSegmentAudios = pendingSegments.map(s => s.audio);
const combinedAudio = Buffer.concat(allSegmentAudios); const combinedAudio = Buffer.concat(allSegmentAudios);
const format = pendingSegments[0].format; const format = pendingSegments[0].format;
// Clear pending segments // Clear pending segments
this.pendingAudioSegments.set(ws, []); session.pendingAudioSegments = [];
this.bufferTimeouts.delete(ws); session.bufferTimeout = null;
// Process combined audio // Process combined audio
await this.processAudio(ws, conversationId, combinedAudio, format); await this.processAudio(ws, conversationId, combinedAudio, format);
} }
}, 10000); // 10 second timeout }, 10000); // 10 second timeout
this.bufferTimeouts.set(ws, timeout);
} }
public close(): void { public close(): void {
this.clients.forEach((_clientId, ws) => { this.sessions.forEach((_session, ws) => {
ws.close(); ws.close();
}); });
this.wss.close(); this.wss.close();