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