From af67235997052b435ab3c2db1a9e50b0991fe8ab Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Mon, 20 Oct 2025 15:28:28 +0200 Subject: [PATCH] refactor: restructure voice assistant server architecture - Introduce a new Session class to encapsulate conversation state and message processing, replacing the previous ClientSession structure. - Implement STTManager and TTSManager for dedicated handling of speech-to-text and text-to-speech functionalities. - Refactor WebSocket server to utilize the new session management, improving message routing and state handling. - Remove deprecated orchestrator functions and streamline message schemas using Zod for validation. - Enhance audio processing with improved buffering and interruption handling, ensuring seamless user interactions. This refactor enhances maintainability and prepares the codebase for future feature expansions. --- .../src/server/agent/orchestrator.ts | 308 +-------- .../src/server/agent/stt-manager.ts | 40 ++ .../src/server/agent/tts-manager.ts | 140 ++-- packages/voice-assistant/src/server/index.ts | 102 +-- .../voice-assistant/src/server/messages.ts | 202 ++++++ .../voice-assistant/src/server/session.ts | 617 ++++++++++++++++++ packages/voice-assistant/src/server/types.ts | 49 -- .../src/server/websocket-server.ts | 439 +++---------- packages/voice-assistant/src/ui/App.tsx | 87 ++- .../src/ui/hooks/useWebSocket.ts | 67 +- 10 files changed, 1114 insertions(+), 937 deletions(-) create mode 100644 packages/voice-assistant/src/server/agent/stt-manager.ts create mode 100644 packages/voice-assistant/src/server/messages.ts create mode 100644 packages/voice-assistant/src/server/session.ts diff --git a/packages/voice-assistant/src/server/agent/orchestrator.ts b/packages/voice-assistant/src/server/agent/orchestrator.ts index e1132861d..a72a719d0 100644 --- a/packages/voice-assistant/src/server/agent/orchestrator.ts +++ b/packages/voice-assistant/src/server/agent/orchestrator.ts @@ -1,303 +1,11 @@ -import { v4 as uuidv4 } from "uuid"; -import { readFile } from "fs/promises"; -import { exec } from "child_process"; -import { promisify } from "util"; -import { getSystemPrompt } from "./system-prompt.js"; -import { streamLLM, type Message } from "./llm-openai.js"; -import { generateTTSAndWaitForPlayback } from "./tts-manager.js"; -import type { VoiceAssistantWebSocketServer } from "../websocket-server.js"; -import type { ArtifactPayload } from "../types.js"; - -const execAsync = promisify(exec); - -interface ConversationContext { - id: string; - messages: Message[]; - createdAt: Date; - lastActivity: Date; -} - /** - * Store active conversations (in-memory for now) - * In production, this could be persisted to a database + * DEPRECATED: This file has been refactored. + * + * All orchestration logic has been moved to session.ts. + * Each Session now owns its own conversation state and manages + * the full lifecycle of user interactions. + * + * This file is kept as a stub to document the migration. + * It can be deleted once the refactoring is confirmed working. */ -const conversations = new Map(); - -/** - * Create a new conversation - */ -export function createConversation(): string { - const id = uuidv4(); - conversations.set(id, { - id, - messages: [], - createdAt: new Date(), - lastActivity: new Date(), - }); - return id; -} - -/** - * Get conversation by ID - */ -export function getConversation(id: string): ConversationContext | null { - return conversations.get(id) || null; -} - -/** - * Delete a conversation by ID - */ -export function deleteConversation(id: string): void { - conversations.delete(id); -} - -/** - * Process user message through the LLM orchestrator - * Handles streaming, tool calls, and WebSocket broadcasting - */ -export async function processUserMessage(params: { - conversationId: string; - message: string; - wsServer?: VoiceAssistantWebSocketServer; - enableTTS?: boolean; - abortSignal?: AbortSignal; -}): Promise { - const conversation = conversations.get(params.conversationId); - if (!conversation) { - throw new Error("Conversation not found"); - } - - // Add user message to context - conversation.messages.push({ - role: "user", - content: params.message, - }); - conversation.lastActivity = new Date(); - - // Note: User message is already broadcast by the caller (e.g., after STT in index.ts) - // No need to broadcast again here to avoid duplication - - let assistantResponse = ""; - // Track pending TTS playback promise outside of streamLLM scope - let pendingTTS: Promise | null = null; - - try { - // Stream LLM response with tool execution - assistantResponse = await streamLLM({ - systemPrompt: getSystemPrompt(), - messages: conversation.messages, - abortSignal: params.abortSignal, - onTextSegment: (segment) => { - // Create TTS promise (don't await it yet) - if (params.wsServer && params.enableTTS) { - pendingTTS = generateTTSAndWaitForPlayback(segment, params.wsServer); - } - - // Broadcast complete text segments - if (params.wsServer) { - params.wsServer.broadcastActivityLog({ - id: uuidv4(), - timestamp: new Date(), - type: "assistant", - content: segment, - }); - } - }, - onChunk: async (chunk) => { - params.wsServer?.broadcast({ - type: "assistant_chunk", - payload: { chunk }, - }); - }, - onToolCall: async (toolCallId, toolName, args) => { - if (pendingTTS) { - console.log("Waiting for pending TTS to finish to execute", toolName); - await pendingTTS; - pendingTTS = null; - } - - // Handle present_artifact tool specially - if (toolName === "present_artifact" && params.wsServer) { - const artifactId = uuidv4(); - - // Resolve source to content - let content: string; - let isBase64 = false; - try { - if (args.source.type === "file") { - const fileBuffer = await readFile(args.source.path); - content = fileBuffer.toString("base64"); - isBase64 = true; - } else if (args.source.type === "command_output") { - const { stdout } = await execAsync(args.source.command, { encoding: 'buffer' }); - content = stdout.toString("base64"); - isBase64 = true; - } else if (args.source.type === "text") { - content = args.source.text; - isBase64 = false; - } else { - content = "[Unknown source type]"; - isBase64 = false; - } - } catch (error) { - console.error("Failed to resolve artifact source:", error); - content = `[Error resolving source: ${error instanceof Error ? error.message : String(error)}]`; - isBase64 = false; - } - - const artifact: ArtifactPayload = { - type: args.type, - id: artifactId, - title: args.title, - content, - isBase64, - }; - - // Broadcast artifact to client - params.wsServer.broadcast({ - type: "artifact", - payload: artifact, - }); - - // Broadcast as activity log entry so it appears in the feed - params.wsServer.broadcastActivityLog({ - id: artifactId, - timestamp: new Date(), - type: "system", - content: `${args.type} artifact: ${args.title}`, - metadata: { artifactId, artifactType: args.type }, - }); - } - - // Broadcast tool call to WebSocket - if (params.wsServer) { - params.wsServer.broadcastActivityLog({ - id: toolCallId, - timestamp: new Date(), - type: "tool_call", - content: `Calling ${toolName}`, - metadata: { toolCallId, toolName, arguments: args }, - }); - } - }, - onToolResult: (toolCallId, toolName, result) => { - // Broadcast tool result to WebSocket - if (params.wsServer) { - params.wsServer.broadcastActivityLog({ - id: toolCallId, - timestamp: new Date(), - type: "tool_result", - content: `Tool ${toolName} completed`, - metadata: { toolCallId, toolName, result }, - }); - } - }, - onToolError: async (toolCallId, toolName, error) => { - // Broadcast tool error to WebSocket - if (params.wsServer) { - params.wsServer.broadcastActivityLog({ - id: toolCallId, - timestamp: new Date(), - type: "error", - content: `Tool ${toolName} failed: ${ - error instanceof Error ? error.message : String(error) - }`, - metadata: { toolCallId, toolName, error }, - }); - } - }, - onError: async (error) => { - // Broadcast general stream error to WebSocket - if (params.wsServer) { - params.wsServer.broadcastActivityLog({ - id: uuidv4(), - timestamp: new Date(), - type: "error", - content: `Stream error: ${ - error instanceof Error ? error.message : String(error) - }`, - }); - } - }, - onFinish: async () => { - // Don't wait for TTS here - we'll handle it after adding to history - }, - }); - - // Add assistant response to context IMMEDIATELY after stream completes - // This ensures partial responses are saved even if TTS fails or is interrupted - conversation.messages.push({ - role: "assistant", - content: assistantResponse, - }); - - // Now wait for any pending TTS, but don't fail the entire operation if it times out - if (pendingTTS) { - try { - await pendingTTS; - } catch (ttsError) { - // TTS failed but message is already in history - just log the error - console.error("TTS playback failed (message already saved):", ttsError); - } - } - } catch (error) { - // If stream itself failed or was aborted, still save any partial response to history - if (assistantResponse) { - conversation.messages.push({ - role: "assistant", - content: assistantResponse, - }); - } - - // Broadcast error to WebSocket - if (params.wsServer) { - params.wsServer.broadcastActivityLog({ - id: uuidv4(), - timestamp: new Date(), - type: "error", - content: `Error: ${ - error instanceof Error ? error.message : String(error) - }`, - }); - } - throw error; - } - - return assistantResponse; -} - -/** - * Clean up old conversations (call periodically) - */ -export function cleanupConversations(maxAgeMinutes: number = 60): void { - const now = new Date(); - for (const [id, conv] of conversations.entries()) { - const ageMinutes = - (now.getTime() - conv.lastActivity.getTime()) / (1000 * 60); - if (ageMinutes > maxAgeMinutes) { - conversations.delete(id); - } - } -} - -/** - * Get conversation statistics - */ -export function getConversationStats(): { - total: number; - conversations: Array<{ - id: string; - messageCount: number; - lastActivity: Date; - }>; -} { - return { - total: conversations.size, - conversations: Array.from(conversations.values()).map((conv) => ({ - id: conv.id, - messageCount: conv.messages.length, - lastActivity: conv.lastActivity, - })), - }; -} diff --git a/packages/voice-assistant/src/server/agent/stt-manager.ts b/packages/voice-assistant/src/server/agent/stt-manager.ts new file mode 100644 index 000000000..a1db5a402 --- /dev/null +++ b/packages/voice-assistant/src/server/agent/stt-manager.ts @@ -0,0 +1,40 @@ +import { transcribeAudio, type TranscriptionResult } from "./stt-openai.js"; + +/** + * Per-session STT manager + * Handles speech-to-text transcription + */ +export class STTManager { + private readonly sessionId: string; + + constructor(sessionId: string) { + this.sessionId = sessionId; + } + + /** + * Transcribe audio buffer to text + */ + public async transcribe( + audio: Buffer, + format: string + ): Promise { + console.log( + `[STT-Manager ${this.sessionId}] Transcribing ${audio.length} bytes of ${format} audio` + ); + + const result = await transcribeAudio(audio, format); + + console.log( + `[STT-Manager ${this.sessionId}] Transcription complete: "${result.text}"` + ); + + return result; + } + + /** + * Cleanup (currently no-op, but provides extension point) + */ + public cleanup(): void { + // No cleanup needed for STT currently + } +} diff --git a/packages/voice-assistant/src/server/agent/tts-manager.ts b/packages/voice-assistant/src/server/agent/tts-manager.ts index 61005c16e..3de932dd1 100644 --- a/packages/voice-assistant/src/server/agent/tts-manager.ts +++ b/packages/voice-assistant/src/server/agent/tts-manager.ts @@ -1,6 +1,6 @@ import { v4 as uuidv4 } from "uuid"; import { synthesizeSpeech } from "./tts-openai.js"; -import type { VoiceAssistantWebSocketServer } from "../websocket-server.js"; +import type { SessionOutboundMessage } from "../messages.js"; interface PendingPlayback { resolve: () => void; @@ -8,68 +8,86 @@ interface PendingPlayback { } /** - * Store pending playback confirmations - * Maps audio ID -> promise resolve/reject handlers + * Per-session TTS manager + * Handles TTS audio generation and playback confirmation tracking */ -const pendingPlaybacks = new Map(); +export class TTSManager { + private pendingPlaybacks: Map = new Map(); + private readonly sessionId: string; -/** - * Generate TTS audio, broadcast to clients, and wait for playback confirmation - * Returns a Promise that resolves when the client confirms playback completed - */ -export async function generateTTSAndWaitForPlayback( - text: string, - wsServer: VoiceAssistantWebSocketServer -): Promise { - // Generate TTS audio - const { audio, format } = await synthesizeSpeech(text); - - // Create unique ID for this audio segment - const audioId = uuidv4(); - - // Create promise that will be resolved when client confirms playback - const playbackPromise = new Promise((resolve, reject) => { - // Store handlers (no timeout - will resolve when client confirms or connection closes) - pendingPlaybacks.set(audioId, { resolve, reject }); - }); - - // Broadcast audio to clients - wsServer.broadcast({ - type: "audio_output", - payload: { - id: audioId, - audio: audio.toString("base64"), - format, - }, - }); - - console.log( - `[TTS-Manager] ${new Date().toISOString()} Sent audio ${audioId}, waiting for playback...` - ); - - // Wait for playback confirmation - await playbackPromise; - - console.log( - `[TTS-Manager] ${new Date().toISOString()} Audio ${audioId} playback confirmed` - ); -} - -/** - * Called when client confirms audio playback completed - * Resolves the corresponding promise - */ -export function confirmAudioPlayed(audioId: string): void { - const pending = pendingPlaybacks.get(audioId); - - if (!pending) { - console.warn( - `[TTS-Manager] Received confirmation for unknown audio ID: ${audioId}` - ); - return; + constructor(sessionId: string) { + this.sessionId = sessionId; } - // Resolve promise and cleanup - pending.resolve(); - pendingPlaybacks.delete(audioId); + /** + * Generate TTS audio, emit to client, and wait for playback confirmation + * Returns a Promise that resolves when the client confirms playback completed + */ + public async generateAndWaitForPlayback( + text: string, + emitMessage: (msg: SessionOutboundMessage) => void + ): Promise { + // Generate TTS audio + const { audio, format } = await synthesizeSpeech(text); + + // Create unique ID for this audio segment + const audioId = uuidv4(); + + // Create promise that will be resolved when client confirms playback + const playbackPromise = new Promise((resolve, reject) => { + // Store handlers (no timeout - will resolve when client confirms or connection closes) + this.pendingPlaybacks.set(audioId, { resolve, reject }); + }); + + // Emit audio output message + emitMessage({ + type: "audio_output", + payload: { + id: audioId, + audio: audio.toString("base64"), + format, + }, + }); + + console.log( + `[TTS-Manager ${this.sessionId}] ${new Date().toISOString()} Sent audio ${audioId}, waiting for playback...` + ); + + // Wait for playback confirmation + await playbackPromise; + + console.log( + `[TTS-Manager ${this.sessionId}] ${new Date().toISOString()} Audio ${audioId} playback confirmed` + ); + } + + /** + * Called when client confirms audio playback completed + * Resolves the corresponding promise + */ + public confirmAudioPlayed(audioId: string): void { + const pending = this.pendingPlaybacks.get(audioId); + + if (!pending) { + console.warn( + `[TTS-Manager ${this.sessionId}] Received confirmation for unknown audio ID: ${audioId}` + ); + return; + } + + // Resolve promise and cleanup + pending.resolve(); + this.pendingPlaybacks.delete(audioId); + } + + /** + * Cleanup all pending playbacks + */ + public cleanup(): void { + // Reject all pending playbacks + for (const [audioId, pending] of this.pendingPlaybacks.entries()) { + pending.reject(new Error("Session closed")); + this.pendingPlaybacks.delete(audioId); + } + } } diff --git a/packages/voice-assistant/src/server/index.ts b/packages/voice-assistant/src/server/index.ts index 93747dd2c..f004094bc 100644 --- a/packages/voice-assistant/src/server/index.ts +++ b/packages/voice-assistant/src/server/index.ts @@ -4,17 +4,12 @@ import path from "path"; import { fileURLToPath } from "url"; import { createServer as createHTTPServer, Server as HttpServer } from "http"; import { readFile } from "fs/promises"; -import { v4 as uuidv4 } from "uuid"; import { createServer as createViteServer } from "vite"; import type { ViteDevServer } from "vite"; import type { ServerConfig } from "./types.js"; import { VoiceAssistantWebSocketServer } from "./websocket-server.js"; -import { initializeSTT, transcribeAudio } from "./agent/stt-openai.js"; +import { initializeSTT } from "./agent/stt-openai.js"; import { initializeTTS } from "./agent/tts-openai.js"; -import { - processUserMessage, - cleanupConversations, -} from "./agent/orchestrator.js"; const __filename = fileURLToPath(import.meta.url); const __dirname = path.dirname(__filename); @@ -132,101 +127,6 @@ async function main() { ); } - // Wire orchestrator to WebSocket for text messages - wsServer.setMessageHandler(async (conversationId: string, message: string, abortSignal: AbortSignal) => { - try { - // Broadcast user's text message as activity log - wsServer.broadcastActivityLog({ - id: uuidv4(), - timestamp: new Date(), - type: "transcript", - content: message, - }); - - await processUserMessage({ - conversationId, - message, - wsServer, - enableTTS: true, - abortSignal, - }); - } catch (error: any) { - console.error("[Orchestrator] Error processing message:", error); - wsServer.broadcastActivityLog({ - id: uuidv4(), - timestamp: new Date(), - type: "error", - content: `Error: ${error.message}`, - }); - } - }); - - // Wire audio handler to WebSocket for voice input (STT) - wsServer.setAudioHandler( - async (conversationId: string, audio: Buffer, format: string, abortSignal: AbortSignal): Promise => { - try { - // Transcribe audio using OpenAI Whisper - const result = await transcribeAudio(audio, format); - - // Check if transcription is empty or only whitespace - const transcriptText = result.text.trim(); - if (!transcriptText) { - console.log("[STT] Transcription is empty or silence, skipping LLM processing"); - // Reset to idle since we're not processing - wsServer.setPhaseForConversation(conversationId, 'idle'); - return ""; - } - - // Broadcast transcription result as activity log - wsServer.broadcastActivityLog({ - id: uuidv4(), - timestamp: new Date(), - type: "transcript", - content: result.text, - metadata: { - language: result.language, - duration: result.duration, - }, - }); - - // Set phase to LLM before processing - wsServer.setPhaseForConversation(conversationId, 'llm'); - - // Process the transcribed text through the orchestrator WITH TTS enabled - // Since this came from voice input, respond with voice output - await processUserMessage({ - conversationId, - message: result.text, - wsServer, - enableTTS: true, - abortSignal, - }); - - // Reset to idle after LLM processing completes - wsServer.setPhaseForConversation(conversationId, 'idle'); - - return result.text; - } catch (error: any) { - // Reset to idle on error - wsServer.setPhaseForConversation(conversationId, 'idle'); - - console.error("[STT] Error transcribing audio:", error); - wsServer.broadcastActivityLog({ - id: uuidv4(), - timestamp: new Date(), - type: "error", - content: `Transcription error: ${error.message}`, - }); - throw error; - } - } - ); - - // Start conversation cleanup interval (every 10 minutes) - setInterval(() => { - cleanupConversations(60); // Clean up conversations older than 60 minutes - }, 10 * 60 * 1000); - httpServer.listen(port, () => { console.log( `\nāœ“ Voice Assistant server running on http://localhost:${port}` diff --git a/packages/voice-assistant/src/server/messages.ts b/packages/voice-assistant/src/server/messages.ts new file mode 100644 index 000000000..76a3e7d66 --- /dev/null +++ b/packages/voice-assistant/src/server/messages.ts @@ -0,0 +1,202 @@ +import { z } from "zod"; + +// ============================================================================ +// Session Inbound Messages (Session receives these) +// ============================================================================ + +export const UserTextMessageSchema = z.object({ + type: z.literal("user_text"), + text: z.string(), +}); + +export const AudioChunkMessageSchema = z.object({ + type: z.literal("audio_chunk"), + audio: z.string(), // base64 encoded + format: z.string(), + isLast: z.boolean(), +}); + +export const AbortRequestMessageSchema = z.object({ + type: z.literal("abort_request"), +}); + +export const AudioPlayedMessageSchema = z.object({ + type: z.literal("audio_played"), + id: z.string(), +}); + +export const SessionInboundMessageSchema = z.discriminatedUnion("type", [ + UserTextMessageSchema, + AudioChunkMessageSchema, + AbortRequestMessageSchema, + AudioPlayedMessageSchema, +]); + +export type SessionInboundMessage = z.infer; + +// ============================================================================ +// Session Outbound Messages (Session emits these) +// ============================================================================ + +export const ActivityLogPayloadSchema = z.object({ + id: z.string(), + timestamp: z.date(), + type: z.enum([ + "transcript", + "assistant", + "tool_call", + "tool_result", + "error", + "system", + ]), + content: z.string(), + metadata: z.record(z.unknown()).optional(), +}); + +export const ActivityLogMessageSchema = z.object({ + type: z.literal("activity_log"), + payload: ActivityLogPayloadSchema, +}); + +export const AssistantChunkMessageSchema = z.object({ + type: z.literal("assistant_chunk"), + payload: z.object({ + chunk: z.string(), + }), +}); + +export const AudioOutputMessageSchema = z.object({ + type: z.literal("audio_output"), + payload: z.object({ + audio: z.string(), // base64 encoded + format: z.string(), + id: z.string(), + }), +}); + +export const TranscriptionResultMessageSchema = z.object({ + type: z.literal("transcription_result"), + payload: z.object({ + text: z.string(), + language: z.string().optional(), + duration: z.number().optional(), + }), +}); + +export const StatusMessageSchema = z.object({ + type: z.literal("status"), + payload: z + .object({ + status: z.string(), + }) + .passthrough(), // Allow additional fields +}); + +export const ArtifactMessageSchema = z.object({ + type: z.literal("artifact"), + payload: z.object({ + type: z.enum(["markdown", "diff", "image", "code"]), + id: z.string(), + title: z.string(), + content: z.string(), + isBase64: z.boolean(), + }), +}); + +export const SessionOutboundMessageSchema = z.discriminatedUnion("type", [ + ActivityLogMessageSchema, + AssistantChunkMessageSchema, + AudioOutputMessageSchema, + TranscriptionResultMessageSchema, + StatusMessageSchema, + ArtifactMessageSchema, +]); + +export type SessionOutboundMessage = z.infer< + typeof SessionOutboundMessageSchema +>; + +// Type exports for individual message types +export type ActivityLogMessage = z.infer; +export type AssistantChunkMessage = z.infer; +export type AudioOutputMessage = z.infer; +export type TranscriptionResultMessage = z.infer; +export type StatusMessage = z.infer; +export type ArtifactMessage = z.infer; + +// Type exports for payload types +export type ActivityLogPayload = z.infer; + +// ============================================================================ +// WebSocket Level Messages (wraps session messages) +// ============================================================================ + +// WebSocket-only messages (not session messages) +export const WSPingMessageSchema = z.object({ + type: z.literal("ping"), +}); + +export const WSPongMessageSchema = z.object({ + type: z.literal("pong"), +}); + +export const WSRecordingStateMessageSchema = z.object({ + type: z.literal("recording_state"), + isRecording: z.boolean(), +}); + +// Wrapped session message +export const WSSessionInboundSchema = z.object({ + type: z.literal("session"), + message: SessionInboundMessageSchema, +}); + +export const WSSessionOutboundSchema = z.object({ + type: z.literal("session"), + message: SessionOutboundMessageSchema, +}); + +// Complete WebSocket message schemas +export const WSInboundMessageSchema = z.discriminatedUnion("type", [ + WSPingMessageSchema, + WSRecordingStateMessageSchema, + WSSessionInboundSchema, +]); + +export const WSOutboundMessageSchema = z.discriminatedUnion("type", [ + WSPongMessageSchema, + WSSessionOutboundSchema, +]); + +export type WSInboundMessage = z.infer; +export type WSOutboundMessage = z.infer; + +// ============================================================================ +// Helper functions for message conversion +// ============================================================================ + +/** + * Extract session message from WebSocket message + * Returns null if message should be handled at WS level only + */ +export function extractSessionMessage( + wsMsg: WSInboundMessage +): SessionInboundMessage | null { + if (wsMsg.type === "session") { + return wsMsg.message; + } + // Ping and recording_state are WS-level only + return null; +} + +/** + * Wrap session message in WebSocket envelope + */ +export function wrapSessionMessage( + sessionMsg: SessionOutboundMessage +): WSOutboundMessage { + return { + type: "session", + message: sessionMsg, + }; +} diff --git a/packages/voice-assistant/src/server/session.ts b/packages/voice-assistant/src/server/session.ts new file mode 100644 index 000000000..2963f6eac --- /dev/null +++ b/packages/voice-assistant/src/server/session.ts @@ -0,0 +1,617 @@ +import { v4 as uuidv4 } from "uuid"; +import { readFile } from "fs/promises"; +import { exec } from "child_process"; +import { promisify } from "util"; +import type { + SessionInboundMessage, + SessionOutboundMessage, +} from "./messages.js"; +import { getSystemPrompt } from "./agent/system-prompt.js"; +import { streamLLM, type Message } from "./agent/llm-openai.js"; +import { TTSManager } from "./agent/tts-manager.js"; +import { STTManager } from "./agent/stt-manager.js"; + +const execAsync = promisify(exec); + +type ProcessingPhase = "idle" | "transcribing" | "llm"; + +/** + * Type for present_artifact tool arguments + */ +interface PresentArtifactArgs { + type: "markdown" | "diff" | "image" | "code"; + source: + | { type: "file"; path: string } + | { type: "command_output"; command: string } + | { type: "text"; text: string }; + title: string; +} + +/** + * Session represents a single client conversation session. + * It owns all state management, orchestration logic, and message processing. + * Session has no knowledge of WebSockets - it only emits and receives messages. + */ +export class Session { + private readonly clientId: string; + private readonly conversationId: string; + private readonly onMessage: (msg: SessionOutboundMessage) => void; + + // State machine + private abortController: AbortController; + private processingPhase: ProcessingPhase = "idle"; + + // Audio buffering for interruption handling + private pendingAudioSegments: Array<{ audio: Buffer; format: string }> = []; + private bufferTimeout: NodeJS.Timeout | null = null; + private audioBuffer: { chunks: Buffer[]; format: string } | null = null; + + // Conversation history + private messages: Message[] = []; + + // Per-session managers + private readonly ttsManager: TTSManager; + private readonly sttManager: STTManager; + + constructor( + clientId: string, + onMessage: (msg: SessionOutboundMessage) => void + ) { + this.clientId = clientId; + this.conversationId = uuidv4(); + this.onMessage = onMessage; + this.abortController = new AbortController(); + + // Initialize per-session managers + this.ttsManager = new TTSManager(this.conversationId); + this.sttManager = new STTManager(this.conversationId); + + console.log( + `[Session ${this.clientId}] Created with conversation ${this.conversationId}` + ); + } + + /** + * Get the conversation ID for this session + */ + public getConversationId(): string { + return this.conversationId; + } + + /** + * Main entry point for processing session messages + */ + public async handleMessage(msg: SessionInboundMessage): Promise { + try { + switch (msg.type) { + case "user_text": + await this.handleUserText(msg.text); + break; + + case "audio_chunk": + await this.handleAudioChunk(msg); + break; + + case "abort_request": + await this.handleAbort(); + break; + + case "audio_played": + this.handleAudioPlayed(msg.id); + break; + } + } catch (error: any) { + console.error(`[Session ${this.clientId}] Error handling message:`, error); + this.emit({ + type: "activity_log", + payload: { + id: uuidv4(), + timestamp: new Date(), + type: "error", + content: `Error: ${error.message}`, + }, + }); + } + } + + /** + * Handle text message from user + */ + private async handleUserText(text: string): Promise { + // Create new abort controller + this.createAbortController(); + + // Emit user message activity log + this.emit({ + type: "activity_log", + payload: { + id: uuidv4(), + timestamp: new Date(), + type: "transcript", + content: text, + }, + }); + + // Add to conversation + this.messages.push({ role: "user", content: text }); + + // Process through LLM (TTS enabled for text input) + await this.processWithLLM(true); + } + + /** + * Handle audio chunk for buffering and transcription + */ + private async handleAudioChunk( + msg: Extract + ): Promise { + // Decode base64 + const audioBuffer = Buffer.from(msg.audio, "base64"); + + if (!msg.isLast) { + // Buffer the chunk + if (!this.audioBuffer) { + this.audioBuffer = { chunks: [], format: msg.format }; + } + this.audioBuffer.chunks.push(audioBuffer); + console.log( + `[Session ${this.clientId}] Buffered audio chunk (${audioBuffer.length} bytes, total: ${this.audioBuffer.chunks.length})` + ); + } else { + // Complete segment received + const allChunks = this.audioBuffer + ? [...this.audioBuffer.chunks, audioBuffer] + : [audioBuffer]; + const format = this.audioBuffer?.format || msg.format; + const currentSegmentAudio = Buffer.concat(allChunks); + + console.log( + `[Session ${this.clientId}] Complete audio segment (${currentSegmentAudio.length} bytes, ${allChunks.length} chunks)` + ); + + // Clear chunk buffer + this.audioBuffer = null; + + // Decision: buffer or process? + const shouldBuffer = + this.processingPhase === "transcribing" && + this.pendingAudioSegments.length === 0; + + if (shouldBuffer) { + // Currently transcribing first segment - buffer this one + console.log( + `[Session ${this.clientId}] Buffering audio segment (phase: ${this.processingPhase})` + ); + this.pendingAudioSegments.push({ + audio: currentSegmentAudio, + format, + }); + this.setBufferTimeout(); + } else if (this.pendingAudioSegments.length > 0) { + // We have buffered segments - add this one and process all together + this.pendingAudioSegments.push({ + audio: currentSegmentAudio, + format, + }); + console.log( + `[Session ${this.clientId}] Processing ${this.pendingAudioSegments.length} buffered segments together` + ); + + // Clear pending segments and timeout + const pendingSegments = [...this.pendingAudioSegments]; + this.pendingAudioSegments = []; + this.clearBufferTimeout(); + + // Concatenate all segments + const allSegmentAudios = pendingSegments.map((s) => s.audio); + const combinedAudio = Buffer.concat(allSegmentAudios); + + await this.processAudio(combinedAudio, format); + } else { + // Normal flow - no buffering needed + await this.processAudio(currentSegmentAudio, format); + } + } + } + + /** + * Process audio through STT and then LLM + */ + private async processAudio(audio: Buffer, format: string): Promise { + this.createAbortController(); + this.setPhase("transcribing"); + + this.emit({ + type: "activity_log", + payload: { + id: uuidv4(), + timestamp: new Date(), + type: "system", + content: "Transcribing audio...", + }, + }); + + try { + const result = await this.sttManager.transcribe(audio, format); + + const transcriptText = result.text.trim(); + if (!transcriptText) { + console.log( + `[Session ${this.clientId}] Empty transcription, skipping LLM` + ); + this.setPhase("idle"); + return; + } + + // Emit transcription result + this.emit({ + type: "transcription_result", + payload: { + text: result.text, + language: result.language, + duration: result.duration, + }, + }); + + // Emit activity log + this.emit({ + type: "activity_log", + payload: { + id: uuidv4(), + timestamp: new Date(), + type: "transcript", + content: result.text, + metadata: { + language: result.language, + duration: result.duration, + }, + }, + }); + + // Add to conversation + this.messages.push({ role: "user", content: result.text }); + + // Set phase to LLM and process + this.setPhase("llm"); + await this.processWithLLM(true); // Enable TTS for voice input + this.setPhase("idle"); + } catch (error: any) { + this.setPhase("idle"); + this.emit({ + type: "activity_log", + payload: { + id: uuidv4(), + timestamp: new Date(), + type: "error", + content: `Transcription error: ${error.message}`, + }, + }); + throw error; + } + } + + /** + * Process user message through LLM with streaming and tool execution + */ + private async processWithLLM(enableTTS: boolean): Promise { + let assistantResponse = ""; + let pendingTTS: Promise | null = null; + + try { + assistantResponse = await streamLLM({ + systemPrompt: getSystemPrompt(), + messages: this.messages, + abortSignal: this.abortController.signal, + onTextSegment: (segment) => { + if (enableTTS) { + // Create TTS promise (don't await yet) + pendingTTS = this.ttsManager.generateAndWaitForPlayback( + segment, + (msg) => this.emit(msg) + ); + } + + // Emit activity log + this.emit({ + type: "activity_log", + payload: { + id: uuidv4(), + timestamp: new Date(), + type: "assistant", + content: segment, + }, + }); + }, + onChunk: async (chunk) => { + this.emit({ + type: "assistant_chunk", + payload: { chunk }, + }); + }, + onToolCall: async (toolCallId, toolName, args) => { + // Wait for pending TTS before executing tool + if (pendingTTS) { + console.log( + `[Session ${this.clientId}] Waiting for TTS before executing ${toolName}` + ); + await pendingTTS; + pendingTTS = null; + } + + // Handle present_artifact tool specially + if (toolName === "present_artifact") { + await this.handlePresentArtifact(toolCallId, args as PresentArtifactArgs); + } + + // Emit tool call activity log + this.emit({ + type: "activity_log", + payload: { + id: toolCallId, + timestamp: new Date(), + type: "tool_call", + content: `Calling ${toolName}`, + metadata: { toolCallId, toolName, arguments: args }, + }, + }); + }, + onToolResult: (toolCallId, toolName, result) => { + this.emit({ + type: "activity_log", + payload: { + id: toolCallId, + timestamp: new Date(), + type: "tool_result", + content: `Tool ${toolName} completed`, + metadata: { toolCallId, toolName, result }, + }, + }); + }, + onToolError: async (toolCallId, toolName, error) => { + this.emit({ + type: "activity_log", + payload: { + id: toolCallId, + timestamp: new Date(), + type: "error", + content: `Tool ${toolName} failed: ${ + error instanceof Error ? error.message : String(error) + }`, + metadata: { toolCallId, toolName, error }, + }, + }); + }, + onError: async (error) => { + this.emit({ + type: "activity_log", + payload: { + id: uuidv4(), + timestamp: new Date(), + type: "error", + content: `Stream error: ${ + error instanceof Error ? error.message : String(error) + }`, + }, + }); + }, + onFinish: async () => { + // Don't wait for TTS here + }, + }); + + // Add assistant response to conversation IMMEDIATELY after stream completes + this.messages.push({ + role: "assistant", + content: assistantResponse, + }); + + // Now wait for any pending TTS, but don't fail if it times out + if (pendingTTS) { + try { + await pendingTTS; + } catch (ttsError) { + console.error( + `[Session ${this.clientId}] TTS playback failed (message already saved):`, + ttsError + ); + } + } + } catch (error) { + // If stream failed or was aborted, still save any partial response + if (assistantResponse) { + this.messages.push({ + role: "assistant", + content: assistantResponse, + }); + } + + this.emit({ + type: "activity_log", + payload: { + id: uuidv4(), + timestamp: new Date(), + type: "error", + content: `Error: ${ + error instanceof Error ? error.message : String(error) + }`, + }, + }); + throw error; + } + } + + /** + * Handle present_artifact tool execution + */ + private async handlePresentArtifact( + toolCallId: string, + args: PresentArtifactArgs + ): Promise { + let content: string; + let isBase64 = false; + + try { + if (args.source.type === "file") { + const fileBuffer = await readFile(args.source.path); + content = fileBuffer.toString("base64"); + isBase64 = true; + } else if (args.source.type === "command_output") { + const { stdout } = await execAsync(args.source.command, { + encoding: "buffer", + }); + content = stdout.toString("base64"); + isBase64 = true; + } else if (args.source.type === "text") { + content = args.source.text; + isBase64 = false; + } else { + content = "[Unknown source type]"; + isBase64 = false; + } + } catch (error) { + console.error( + `[Session ${this.clientId}] Failed to resolve artifact source:`, + error + ); + content = `[Error: ${error instanceof Error ? error.message : String(error)}]`; + isBase64 = false; + } + + // Emit artifact message + this.emit({ + type: "artifact", + payload: { + type: args.type, + id: toolCallId, + title: args.title, + content, + isBase64, + }, + }); + + // Emit activity log for artifact + this.emit({ + type: "activity_log", + payload: { + id: toolCallId, + timestamp: new Date(), + type: "system", + content: `${args.type} artifact: ${args.title}`, + metadata: { artifactId: toolCallId, artifactType: args.type }, + }, + }); + } + + /** + * Handle abort request from client + */ + private async handleAbort(): Promise { + console.log( + `[Session ${this.clientId}] Abort request, phase: ${this.processingPhase}` + ); + + if (this.processingPhase === "llm") { + // Already in LLM phase - abort immediately + this.abortController.abort(); + console.log(`[Session ${this.clientId}] Aborted LLM processing`); + + // Reset phase to idle + this.setPhase("idle"); + + // Clear any pending segments and timeouts + this.pendingAudioSegments = []; + this.clearBufferTimeout(); + } else if (this.processingPhase === "transcribing") { + // Still in STT phase - we'll buffer the next audio + console.log( + `[Session ${this.clientId}] Will buffer next audio (currently transcribing)` + ); + // Phase stays as 'transcribing', handleAudioChunk will handle buffering + } + // If idle, nothing to do + } + + /** + * Handle audio playback confirmation from client + */ + private handleAudioPlayed(id: string): void { + this.ttsManager.confirmAudioPlayed(id); + } + + /** + * Create new AbortController, aborting the previous one + */ + private createAbortController(): AbortController { + this.abortController.abort(); + this.abortController = new AbortController(); + return this.abortController; + } + + /** + * Set the processing phase + */ + private setPhase(phase: ProcessingPhase): void { + this.processingPhase = phase; + console.log(`[Session ${this.clientId}] Phase: ${phase}`); + } + + /** + * Set timeout to process buffered audio segments + */ + private setBufferTimeout(): void { + this.clearBufferTimeout(); + + this.bufferTimeout = setTimeout(async () => { + console.log( + `[Session ${this.clientId}] Buffer timeout reached, processing pending segments` + ); + + if (this.pendingAudioSegments.length > 0) { + const segments = [...this.pendingAudioSegments]; + this.pendingAudioSegments = []; + this.bufferTimeout = null; + + const combined = Buffer.concat(segments.map((s) => s.audio)); + await this.processAudio(combined, segments[0].format); + } + }, 10000); // 10 second timeout + } + + /** + * Clear buffer timeout + */ + private clearBufferTimeout(): void { + if (this.bufferTimeout) { + clearTimeout(this.bufferTimeout); + this.bufferTimeout = null; + } + } + + /** + * Emit a message to the client + */ + private emit(msg: SessionOutboundMessage): void { + this.onMessage(msg); + } + + /** + * Clean up session resources + */ + public cleanup(): void { + console.log(`[Session ${this.clientId}] Cleaning up`); + + // Abort any ongoing operations + this.abortController.abort(); + + // Clear timeouts + this.clearBufferTimeout(); + + // Clear buffers + this.pendingAudioSegments = []; + this.audioBuffer = null; + + // Cleanup managers + this.ttsManager.cleanup(); + this.sttManager.cleanup(); + } +} diff --git a/packages/voice-assistant/src/server/types.ts b/packages/voice-assistant/src/server/types.ts index 2e07e2f2e..4bb2190d5 100644 --- a/packages/voice-assistant/src/server/types.ts +++ b/packages/voice-assistant/src/server/types.ts @@ -4,52 +4,3 @@ export interface ServerConfig { port: number; isDev: boolean; } - -export interface ActivityLogEntry { - id: string; - timestamp: Date; - type: 'transcript' | 'assistant' | 'tool_call' | 'tool_result' | 'error' | 'system'; - content: string; - metadata?: Record; -} - -export interface WebSocketMessage { - type: 'activity_log' | 'status' | 'ping' | 'pong' | 'user_message' | 'assistant_chunk' - | 'audio_chunk' | 'audio_output' | 'recording_state' | 'transcription_result' | 'audio_played' - | 'artifact' | 'abort_request'; - payload: unknown; -} - -export interface AudioChunkPayload { - audio: string; // base64 encoded audio data - format: string; // 'webm', 'ogg', etc. - isLast: boolean; // true when recording stopped -} - -export interface RecordingStatePayload { - isRecording: boolean; -} - -export interface TranscriptionResultPayload { - text: string; - language?: string; - duration?: number; -} - -export interface AudioOutputPayload { - audio: string; // base64 encoded audio data (complete) - format: string; // 'mp3' - id: string; // unique ID for queue management -} - -export interface AudioPlayedPayload { - id: string; // unique ID of the audio that finished playing -} - -export interface ArtifactPayload { - type: 'markdown' | 'diff' | 'image' | 'code'; - id: string; - title: string; - content: string; // resolved content (base64 for file/command sources, plain text for text source) - isBase64: boolean; // true if content is base64 encoded (from file or command_output) -} diff --git a/packages/voice-assistant/src/server/websocket-server.ts b/packages/voice-assistant/src/server/websocket-server.ts index 4abac1e00..a97705199 100644 --- a/packages/voice-assistant/src/server/websocket-server.ts +++ b/packages/voice-assistant/src/server/websocket-server.ts @@ -1,76 +1,21 @@ import { WebSocketServer, WebSocket } from "ws"; import { Server as HTTPServer } from "http"; -import type { - WebSocketMessage, - ActivityLogEntry, - AudioChunkPayload, - AudioPlayedPayload, -} from "./types.js"; -import { confirmAudioPlayed } from "./agent/tts-manager.js"; -import { createConversation, deleteConversation } from "./agent/orchestrator.js"; - -type ProcessingPhase = 'idle' | 'transcribing' | 'llm'; +import { + WSInboundMessageSchema, + type WSOutboundMessage, + extractSessionMessage, + wrapSessionMessage, +} from "./messages.js"; +import { Session } from "./session.js"; /** - * Encapsulates all state for a single WebSocket client session + * WebSocket server that routes messages between clients and their sessions. + * This is a thin transport layer with no business logic. */ -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 sessions: Map = new Map(); - private conversationIdToWs: Map = new Map(); // Reverse map for phase management - private messageHandler?: (conversationId: string, message: string, abortSignal: AbortSignal) => Promise; - private audioHandler?: (conversationId: string, audio: Buffer, format: string, abortSignal: AbortSignal) => Promise; + private sessions: Map = new Map(); + private conversationIdToWs: Map = new Map(); private clientIdCounter: number = 0; constructor(server: HTTPServer) { @@ -83,159 +28,140 @@ export class VoiceAssistantWebSocketServer { console.log("āœ“ WebSocket server initialized on /ws"); } + /** + * Handle new WebSocket connection + */ private handleConnection(ws: WebSocket): void { - // Generate unique client ID and conversation + // Generate unique client ID const clientId = `client-${++this.clientIdCounter}`; - const conversationId = createConversation(); - // Create session for this client - const session = new ClientSession(clientId, conversationId); + // Create session with message emission callback + const session = new Session(clientId, (msg) => { + this.sendToClient(ws, wrapSessionMessage(msg)); + }); + + // Store session and reverse mapping this.sessions.set(ws, session); - this.conversationIdToWs.set(conversationId, ws); + this.conversationIdToWs.set(session.getConversationId(), ws); console.log( - `[WS] Client connected: ${clientId} with conversation ${conversationId} (total: ${this.sessions.size})` + `[WS] Client connected: ${clientId} with conversation ${session.getConversationId()} (total: ${this.sessions.size})` ); // Send welcome message - this.sendToClient(ws, { + this.sendToClient(ws, wrapSessionMessage({ type: "status", payload: { status: "connected", message: "WebSocket connection established", }, - }); + })); + // Set up message handler ws.on("message", (data) => { this.handleMessage(ws, data); }); + // Set up close handler ws.on("close", () => { const session = this.sessions.get(ws); if (!session) return; - console.log(`[WS] Client disconnected: ${session.clientId} (total: ${this.sessions.size - 1})`); + console.log( + `[WS] Client disconnected: ${clientId} (total: ${this.sessions.size - 1})` + ); // Clean up session session.cleanup(); - deleteConversation(session.conversationId); // Remove from maps this.sessions.delete(ws); - this.conversationIdToWs.delete(session.conversationId); + this.conversationIdToWs.delete(session.getConversationId()); - console.log(`[WS] Conversation ${session.conversationId} deleted`); + console.log(`[WS] Conversation ${session.getConversationId()} deleted`); }); + // Set up error handler ws.on("error", (error) => { - console.error("[WS] Client error:", error); + console.error(`[WS] Client error:`, error); const session = this.sessions.get(ws); if (!session) return; // Clean up session session.cleanup(); - deleteConversation(session.conversationId); // Remove from maps this.sessions.delete(ws); - this.conversationIdToWs.delete(session.conversationId); + this.conversationIdToWs.delete(session.getConversationId()); }); } + /** + * Handle incoming WebSocket message + */ private async handleMessage( ws: WebSocket, data: Buffer | ArrayBuffer | Buffer[] ): Promise { try { - const message = JSON.parse(data.toString()) as WebSocketMessage; + // Parse message + const parsed = JSON.parse(data.toString()); + + // Validate with Zod + const message = WSInboundMessageSchema.parse(parsed); console.log(`[WS] Received message type: ${message.type}`); + // Handle WebSocket-level messages switch (message.type) { case "ping": - this.sendToClient(ws, { type: "pong", payload: {} }); - break; + this.sendToClient(ws, { type: "pong" }); + return; - case "user_message": - // Handle user message through orchestrator - const payload = message.payload as { message: string }; - const session = this.sessions.get(ws); - if (!session) { - console.error("[WS] No session found for client"); - break; + case "recording_state": + console.log(`[WS] Recording state: ${message.isRecording}`); + return; + + case "session": + // Extract and forward session message + const sessionMessage = extractSessionMessage(message); + if (sessionMessage) { + const session = this.sessions.get(ws); + if (session) { + await session.handleMessage(sessionMessage); + } else { + console.error("[WS] No session found for client"); + } } - - // Create new abort controller for this request (aborts previous one) - const newController = session.createAbortController(); - - if (this.messageHandler) { - await this.messageHandler(session.conversationId, payload.message, newController.signal); - } else { - console.warn("[WS] No message handler registered"); - } - break; - - case "audio_chunk": - // Handle audio chunk for STT - await this.handleAudioChunk(ws, message.payload as AudioChunkPayload); - break; - - case "audio_played": - // Handle audio playback confirmation - const audioPlayedPayload = message.payload as AudioPlayedPayload; - confirmAudioPlayed(audioPlayedPayload.id); - break; - - case "abort_request": - // Handle abort request from client (e.g., when VAD detects new speech) - await this.handleAbortRequest(ws); - break; - - default: - console.warn(`[WS] Unknown message type: ${message.type}`); + return; } - } catch (error) { - console.error("[WS] Failed to parse message:", error); + } catch (error: any) { + console.error("[WS] Failed to parse/handle message:", error); + // Send error to client + this.sendToClient(ws, wrapSessionMessage({ + type: "status", + payload: { + status: "error", + message: `Invalid message: ${error.message}`, + }, + })); } } - private async handleAbortRequest(ws: WebSocket): Promise { - const session = this.sessions.get(ws); - if (!session) return; - - console.log(`[WS] Abort request received, current phase: ${session.processingPhase}`); - - if (session.processingPhase === 'llm') { - // Already in LLM phase - abort immediately - session.abortController.abort(); - console.log(`[WS] Aborted LLM processing`); - - // Reset phase to idle - session.setPhase('idle'); - - // Clear any pending segments and timeouts - session.pendingAudioSegments = []; - if (session.bufferTimeout) { - clearTimeout(session.bufferTimeout); - session.bufferTimeout = null; - } - } 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)`); - // Phase stays as 'transcribing', handleAudioChunk will handle buffering - } - // If idle, nothing to do - } - - private sendToClient(ws: WebSocket, message: WebSocketMessage): void { + /** + * Send message to specific client + */ + private sendToClient(ws: WebSocket, message: WSOutboundMessage): void { if (ws.readyState === 1) { // WebSocket.OPEN = 1 ws.send(JSON.stringify(message)); } } - public broadcast(message: WebSocketMessage): void { + /** + * Broadcast message to all connected clients + */ + public broadcast(message: WSOutboundMessage): void { const payload = JSON.stringify(message); this.sessions.forEach((_session, client) => { if (client.readyState === 1) { @@ -245,209 +171,12 @@ export class VoiceAssistantWebSocketServer { }); } - public broadcastActivityLog(entry: ActivityLogEntry): void { - this.broadcast({ - type: "activity_log", - payload: entry, - }); - } - - public broadcastStatus( - status: string, - metadata?: Record - ): void { - this.broadcast({ - type: "status", - payload: { status, ...metadata }, - }); - } - - public setMessageHandler(handler: (conversationId: string, message: string, abortSignal: AbortSignal) => Promise): void { - this.messageHandler = handler; - } - - public setAudioHandler( - handler: (conversationId: string, audio: Buffer, format: string, abortSignal: AbortSignal) => Promise - ): void { - this.audioHandler = handler; - } - - public setPhaseForConversation(conversationId: string, phase: ProcessingPhase): void { - const ws = this.conversationIdToWs.get(conversationId); - if (ws) { - const session = this.sessions.get(ws); - if (session) { - session.setPhase(phase); - } - } else { - console.warn(`[WS] Cannot set phase for unknown conversation ${conversationId}`); - } - } - - private async handleAudioChunk( - ws: WebSocket, - payload: AudioChunkPayload - ): Promise { - try { - const session = this.sessions.get(ws); - if (!session) { - console.error("[WS] No session found for WebSocket"); - return; - } - - // Decode base64 audio data - const audioBuffer = Buffer.from(payload.audio, "base64"); - - if (!payload.isLast) { - // Buffer the chunk - if (!session.audioBuffer) { - session.audioBuffer = { - chunks: [], - format: payload.format, - }; - } - session.audioBuffer.chunks.push(audioBuffer); - console.log( - `[WS] Buffered audio chunk (${audioBuffer.length} bytes, total chunks: ${session.audioBuffer.chunks.length})` - ); - } else { - // Last chunk - complete audio segment received - const allChunks = session.audioBuffer - ? [...session.audioBuffer.chunks, audioBuffer] - : [audioBuffer]; - const format = session.audioBuffer?.format || payload.format; - - // Concatenate all chunks for this segment - const currentSegmentAudio = Buffer.concat(allChunks); - console.log( - `[WS] Complete audio segment received (${currentSegmentAudio.length} bytes, ${allChunks.length} chunks)` - ); - - // Clear chunk buffer - session.audioBuffer = null; - - // Decision: buffer or process? - 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: ${session.processingPhase})`); - session.pendingAudioSegments.push({ audio: currentSegmentAudio, format }); - - // Set timeout to process buffer if no more audio arrives - this.setBufferTimeout(ws, session.conversationId); - } else if (session.pendingAudioSegments.length > 0) { - // We have buffered segments - add this one and process all together - session.pendingAudioSegments.push({ audio: currentSegmentAudio, format }); - console.log(`[WS] Processing ${session.pendingAudioSegments.length} buffered audio segments together`); - - // Clear pending segments and timeout - const pendingSegments = [...session.pendingAudioSegments]; - session.pendingAudioSegments = []; - if (session.bufferTimeout) { - clearTimeout(session.bufferTimeout); - session.bufferTimeout = null; - } - - // Concatenate all segments - const allSegmentAudios = pendingSegments.map(s => s.audio); - const combinedAudio = Buffer.concat(allSegmentAudios); - - // Process combined audio - await this.processAudio(ws, session.conversationId, combinedAudio, format); - } else { - // Normal flow - no buffering needed - await this.processAudio(ws, session.conversationId, currentSegmentAudio, format); - } - } - } catch (error: any) { - console.error("[WS] Audio chunk handling error:", error); - this.broadcastActivityLog({ - id: Date.now().toString(), - timestamp: new Date(), - type: "error", - content: `Audio processing error: ${error.message}`, - }); - } - } - - private async processAudio( - ws: WebSocket, - conversationId: string, - audio: Buffer, - format: string - ): Promise { - const session = this.sessions.get(ws); - if (!session) return; - - // Create new abort controller for this request (aborts previous one) - const newAbortController = session.createAbortController(); - - // Set phase to transcribing - session.setPhase('transcribing'); - - if (this.audioHandler) { - this.broadcastActivityLog({ - id: Date.now().toString(), - timestamp: new Date(), - type: "system", - content: "Transcribing audio...", - }); - - try { - const transcript = await this.audioHandler(conversationId, audio, format, newAbortController.signal); - - // Send transcription result back to client - this.sendToClient(ws, { - type: "transcription_result", - payload: { text: transcript }, - }); - - // Phase management is now handled by the audioHandler in index.ts - // 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 - session.setPhase('idle'); - throw error; - } - } else { - console.warn("[WS] No audio handler registered"); - session.setPhase('idle'); - } - } - - private setBufferTimeout(ws: WebSocket, conversationId: string): void { - const session = this.sessions.get(ws); - if (!session) return; - - // Clear any existing timeout - if (session.bufferTimeout) { - clearTimeout(session.bufferTimeout); - } - - // Set new timeout (10 seconds) - session.bufferTimeout = setTimeout(async () => { - console.log(`[WS] Buffer timeout reached, processing pending segments`); - - 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 - session.pendingAudioSegments = []; - session.bufferTimeout = null; - - // Process combined audio - await this.processAudio(ws, conversationId, combinedAudio, format); - } - }, 10000); // 10 second timeout - } - + /** + * Close the WebSocket server + */ public close(): void { - this.sessions.forEach((_session, ws) => { + this.sessions.forEach((session, ws) => { + session.cleanup(); ws.close(); }); this.wss.close(); diff --git a/packages/voice-assistant/src/ui/App.tsx b/packages/voice-assistant/src/ui/App.tsx index a6ec9ab3a..4b46ec23e 100644 --- a/packages/voice-assistant/src/ui/App.tsx +++ b/packages/voice-assistant/src/ui/App.tsx @@ -105,31 +105,23 @@ function App() { useEffect(() => { // Listen for status messages - const unsubStatus = ws.on("status", (payload: unknown) => { - const data = payload as { status: string; message?: string }; - addLog("info", data.message || `Status: ${data.status}`); - }); - - // Listen for pong messages - const unsubPong = ws.on("pong", () => { - addLog("success", "Received pong from server"); + const unsubStatus = ws.on("status", (message) => { + if (message.type !== 'status') return; + const msg = 'message' in message.payload ? String(message.payload.message) : undefined; + addLog("info", msg || `Status: ${message.payload.status}`); }); // Listen for activity log messages - const unsubActivity = ws.on("activity_log", (payload: unknown) => { - const data = payload as { - id: string; - type: string; - content: string; - metadata?: Record; - }; + const unsubActivity = ws.on("activity_log", (message) => { + if (message.type !== 'activity_log') return; + const data = message.payload; // Handle tool calls - if (data.type === "tool_call") { + if (data.type === "tool_call" && data.metadata) { const { toolCallId, toolName, arguments: args } = data.metadata as { toolCallId: string; toolName: string; - arguments: any; + arguments: unknown; }; setLogs((prev) => [ @@ -146,10 +138,10 @@ function App() { return; } - if (data.type === "tool_result") { + if (data.type === "tool_result" && data.metadata) { const { toolCallId, result } = data.metadata as { toolCallId: string; - result: any; + result: unknown; }; setLogs((prev) => @@ -163,10 +155,10 @@ function App() { } // Handle tool errors - update tool call status to failed - if (data.type === "error" && data.metadata?.toolCallId) { + if (data.type === "error" && data.metadata && 'toolCallId' in data.metadata) { const { toolCallId, error } = data.metadata as { toolCallId: string; - error: any; + error: unknown; }; setLogs((prev) => @@ -194,24 +186,22 @@ function App() { }); // Listen for streaming assistant chunks - const unsubChunk = ws.on("assistant_chunk", (payload: unknown) => { - const data = payload as { chunk: string }; - setCurrentAssistantMessage((prev) => prev + data.chunk); + const unsubChunk = ws.on("assistant_chunk", (message) => { + if (message.type !== 'assistant_chunk') return; + setCurrentAssistantMessage((prev) => prev + message.payload.chunk); }); // Listen for transcription results - const unsubTranscription = ws.on( - "transcription_result", - (_payload: unknown) => { - // Note: Transcription is already broadcast as activity_log with type "transcript" - // No need to log it again here to avoid duplication - setIsProcessingAudio(false); - } - ); + const unsubTranscription = ws.on("transcription_result", () => { + // Note: Transcription is already broadcast as activity_log with type "transcript" + // No need to log it again here to avoid duplication + setIsProcessingAudio(false); + }); // Listen for artifacts - const unsubArtifact = ws.on("artifact", (payload: unknown) => { - const artifact = payload as Artifact; + const unsubArtifact = ws.on("artifact", (message) => { + if (message.type !== 'artifact') return; + const artifact = message.payload; // Add to artifacts map setArtifacts((prev) => { @@ -238,8 +228,9 @@ function App() { }); // Listen for audio output (TTS) - const unsubAudioOutput = ws.on("audio_output", async (payload: unknown) => { - const data = payload as { audio: string; format: string; id: string }; + const unsubAudioOutput = ws.on("audio_output", async (message) => { + if (message.type !== 'audio_output') return; + const data = message.payload; try { setIsPlayingAudio(true); @@ -261,8 +252,11 @@ function App() { // Send confirmation back to server ws.send({ - type: "audio_played", - payload: { id: data.id }, + type: "session", + message: { + type: "audio_played", + id: data.id, + }, }); setIsPlayingAudio(false); @@ -275,7 +269,6 @@ function App() { return () => { unsubStatus(); - unsubPong(); unsubActivity(); unsubChunk(); unsubTranscription(); @@ -337,8 +330,9 @@ function App() { ); ws.send({ - type: "audio_chunk", - payload: { + type: "session", + message: { + type: "audio_chunk", audio: base64Audio, format: format, isLast: true, @@ -438,8 +432,10 @@ function App() { // Send abort request to server immediately ws.send({ - type: "abort_request", - payload: {}, + type: "session", + message: { + type: "abort_request", + }, }); }, onSpeechEnd: async (audioData: Float32Array) => { @@ -462,8 +458,9 @@ function App() { ); ws.send({ - type: "audio_chunk", - payload: { + type: "session", + message: { + type: "audio_chunk", audio: base64Audio, format: audioBlob.type, isLast: true, diff --git a/packages/voice-assistant/src/ui/hooks/useWebSocket.ts b/packages/voice-assistant/src/ui/hooks/useWebSocket.ts index ccedb21da..6569f61e6 100644 --- a/packages/voice-assistant/src/ui/hooks/useWebSocket.ts +++ b/packages/voice-assistant/src/ui/hooks/useWebSocket.ts @@ -1,16 +1,14 @@ import { useEffect, useRef, useState, useCallback } from 'react'; - -interface WebSocketMessage { - type: string; - payload: unknown; -} - -type MessageHandler = (payload: unknown) => void; +import type { + WSInboundMessage, + WSOutboundMessage, + SessionOutboundMessage +} from '../../server/messages.js'; export interface UseWebSocketReturn { isConnected: boolean; - send: (message: WebSocketMessage) => void; - on: (type: string, handler: MessageHandler) => () => void; + send: (message: WSInboundMessage) => void; + on: (type: SessionOutboundMessage['type'], handler: (message: SessionOutboundMessage) => void) => () => void; sendPing: () => void; sendUserMessage: (message: string) => void; } @@ -18,7 +16,7 @@ export interface UseWebSocketReturn { export function useWebSocket(url: string): UseWebSocketReturn { const [isConnected, setIsConnected] = useState(false); const wsRef = useRef(null); - const handlersRef = useRef>>(new Map()); + const handlersRef = useRef void>>>(new Map()); const reconnectTimeoutRef = useRef(); const connect = useCallback(() => { @@ -51,19 +49,27 @@ export function useWebSocket(url: string): UseWebSocketReturn { ws.onmessage = (event) => { try { - const message = JSON.parse(event.data) as WebSocketMessage; - console.log(`[WS] Received message type: ${message.type}`, message.payload); + const wsMessage: WSOutboundMessage = JSON.parse(event.data); - // Call all registered handlers for this message type - const handlers = handlersRef.current.get(message.type); - if (handlers) { - handlers.forEach((handler) => { - try { - handler(message.payload); - } catch (err) { - console.error(`[WS] Error in handler for ${message.type}:`, err); - } - }); + // Only session messages trigger handlers + if (wsMessage.type === 'session') { + const sessionMessage = wsMessage.message; + console.log(`[WS] Received session message type: ${sessionMessage.type}`); + + // Call all registered handlers for this message type + const handlers = handlersRef.current.get(sessionMessage.type); + if (handlers) { + handlers.forEach((handler) => { + try { + handler(sessionMessage); + } catch (err) { + console.error(`[WS] Error in handler for ${sessionMessage.type}:`, err); + } + }); + } + } else { + // pong - just log + console.log(`[WS] Received ${wsMessage.type}`); } } catch (err) { console.error('[WS] Failed to parse message:', err); @@ -89,7 +95,7 @@ export function useWebSocket(url: string): UseWebSocketReturn { }; }, [connect]); - const send = useCallback((message: WebSocketMessage) => { + const send = useCallback((message: WSInboundMessage) => { if (wsRef.current?.readyState === WebSocket.OPEN) { wsRef.current.send(JSON.stringify(message)); } else { @@ -97,7 +103,10 @@ export function useWebSocket(url: string): UseWebSocketReturn { } }, []); - const on = useCallback((type: string, handler: MessageHandler) => { + const on = useCallback(( + type: SessionOutboundMessage['type'], + handler: (message: SessionOutboundMessage) => void + ) => { if (!handlersRef.current.has(type)) { handlersRef.current.set(type, new Set()); } @@ -116,12 +125,18 @@ export function useWebSocket(url: string): UseWebSocketReturn { }, []); const sendPing = useCallback(() => { - send({ type: 'ping', payload: {} }); + send({ type: 'ping' }); }, [send]); const sendUserMessage = useCallback( (message: string) => { - send({ type: 'user_message', payload: { message } }); + send({ + type: 'session', + message: { + type: 'user_text', + text: message, + }, + }); }, [send] );