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.
2025-10-20 15:28:28 +02:00
parent 8365fb889b
commit af67235997
10 changed files with 1114 additions and 937 deletions
@@ -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)
+ * DEPRECATED: This file has been refactored.
- * In production, this could be persisted to a database
+ *
 * 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<string, ConversationContext>();
 /**
 * 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<string> {
  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<void> | 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,
    })),
  };
 }
@@ -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<TranscriptionResult> {
    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
  }
 }
@@ -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
+ * Per-session TTS manager
- * Maps audio ID -> promise resolve/reject handlers
+ * Handles TTS audio generation and playback confirmation tracking
 */
-const pendingPlaybacks = new Map<string, PendingPlayback>();
+export class TTSManager {
  private pendingPlaybacks: Map<string, PendingPlayback> = new Map();
  private readonly sessionId: string;
-/**
+  constructor(sessionId: string) {
- * Generate TTS audio, broadcast to clients, and wait for playback confirmation
+    this.sessionId = sessionId;
 * Returns a Promise that resolves when the client confirms playback completed
 */
 export async function generateTTSAndWaitForPlayback(
  text: string,
  wsServer: VoiceAssistantWebSocketServer
 ): Promise<void> {
  // 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<void>((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;
  }
-  // Resolve promise and cleanup
+  /**
-  pending.resolve();
+   * Generate TTS audio, emit to client, and wait for playback confirmation
-  pendingPlaybacks.delete(audioId);
+   * Returns a Promise that resolves when the client confirms playback completed
   */
  public async generateAndWaitForPlayback(
    text: string,
    emitMessage: (msg: SessionOutboundMessage) => void
  ): Promise<void> {
    // 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<void>((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);
    }
  }
 }
@@ -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<string> => {
      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}`
@@ -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<typeof SessionInboundMessageSchema>;
 // ============================================================================
 // 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<typeof ActivityLogMessageSchema>;
 export type AssistantChunkMessage = z.infer<typeof AssistantChunkMessageSchema>;
 export type AudioOutputMessage = z.infer<typeof AudioOutputMessageSchema>;
 export type TranscriptionResultMessage = z.infer<typeof TranscriptionResultMessageSchema>;
 export type StatusMessage = z.infer<typeof StatusMessageSchema>;
 export type ArtifactMessage = z.infer<typeof ArtifactMessageSchema>;
 // Type exports for payload types
 export type ActivityLogPayload = z.infer<typeof ActivityLogPayloadSchema>;
 // ============================================================================
 // 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<typeof WSInboundMessageSchema>;
 export type WSOutboundMessage = z.infer<typeof WSOutboundMessageSchema>;
 // ============================================================================
 // 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,
  };
 }
@@ -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<void> {
    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<void> {
    // 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<SessionInboundMessage, { type: "audio_chunk" }>
  ): Promise<void> {
    // 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<void> {
    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<void> {
    let assistantResponse = "";
    let pendingTTS: Promise<void> | 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<void> {
    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<void> {
    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();
  }
 }
@@ -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<string, unknown>;
 }
 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)
 }
@@ -1,76 +1,21 @@
 import { WebSocketServer, WebSocket } from "ws";
 import { Server as HTTPServer } from "http";
-import type {
+import {
-  WebSocketMessage,
+  WSInboundMessageSchema,
-  ActivityLogEntry,
+  type WSOutboundMessage,
-  AudioChunkPayload,
+  extractSessionMessage,
-  AudioPlayedPayload,
+  wrapSessionMessage,
-} from "./types.js";
+} from "./messages.js";
-import { confirmAudioPlayed } from "./agent/tts-manager.js";
+import { Session } from "./session.js";
 import { createConversation, deleteConversation } from "./agent/orchestrator.js";
 type ProcessingPhase = 'idle' | 'transcribing' | 'llm';
 /**
- * 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<WebSocket, ClientSession> = new Map();
+  private sessions: Map<WebSocket, Session> = new Map();
-  private conversationIdToWs: Map<string, WebSocket> = new Map(); // Reverse map for phase management
+  private conversationIdToWs: Map<string, WebSocket> = new Map();
  private messageHandler?: (conversationId: string, message: string, abortSignal: AbortSignal) => Promise<void>;
  private audioHandler?: (conversationId: string, audio: Buffer, format: string, abortSignal: AbortSignal) => Promise<string>;
  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
+    // Create session with message emission callback
-    const session = new ClientSession(clientId, conversationId);
+    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<void> {
    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: {} });
+          this.sendToClient(ws, { type: "pong" });
-          break;
+          return;
-        case "user_message":
+        case "recording_state":
-          // Handle user message through orchestrator
+          console.log(`[WS] Recording state: ${message.isRecording}`);
-          const payload = message.payload as { message: string };
+          return;
-          const session = this.sessions.get(ws);
+
-          if (!session) {
+        case "session":
-            console.error("[WS] No session found for client");
+          // Extract and forward session message
-            break;
+          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");
            }
          }
-
+          return;
          // 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}`);
      }
-    } catch (error) {
+    } catch (error: any) {
-      console.error("[WS] Failed to parse message:", error);
+      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<void> {
+  /**
-    const session = this.sessions.get(ws);
+   * Send message to specific client
-    if (!session) return;
+   */
-
+  private sendToClient(ws: WebSocket, message: WSOutboundMessage): void {
    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 {
    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({
+   * Close the WebSocket server
-      type: "activity_log",
+   */
      payload: entry,
    });
  }
  public broadcastStatus(
    status: string,
    metadata?: Record<string, unknown>
  ): void {
    this.broadcast({
      type: "status",
      payload: { status, ...metadata },
    });
  }
  public setMessageHandler(handler: (conversationId: string, message: string, abortSignal: AbortSignal) => Promise<void>): void {
    this.messageHandler = handler;
  }
  public setAudioHandler(
    handler: (conversationId: string, audio: Buffer, format: string, abortSignal: AbortSignal) => Promise<string>
  ): 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<void> {
    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<void> {
    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
  }
  public close(): void {
-    this.sessions.forEach((_session, ws) => {
+    this.sessions.forEach((session, ws) => {
      session.cleanup();
      ws.close();
    });
    this.wss.close();
@@ -105,31 +105,23 @@ function App() {
  useEffect(() => {
    // Listen for status messages
-    const unsubStatus = ws.on("status", (payload: unknown) => {
+    const unsubStatus = ws.on("status", (message) => {
-      const data = payload as { status: string; message?: string };
+      if (message.type !== 'status') return;
-      addLog("info", data.message || `Status: ${data.status}`);
+      const msg = 'message' in message.payload ? String(message.payload.message) : undefined;
-    });
+      addLog("info", msg || `Status: ${message.payload.status}`);
    // Listen for pong messages
    const unsubPong = ws.on("pong", () => {
      addLog("success", "Received pong from server");
    });
    // Listen for activity log messages
-    const unsubActivity = ws.on("activity_log", (payload: unknown) => {
+    const unsubActivity = ws.on("activity_log", (message) => {
-      const data = payload as {
+      if (message.type !== 'activity_log') return;
-        id: string;
+      const data = message.payload;
        type: string;
        content: string;
        metadata?: Record<string, unknown>;
      };
      // 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 unsubChunk = ws.on("assistant_chunk", (message) => {
-      const data = payload as { chunk: string };
+      if (message.type !== 'assistant_chunk') return;
-      setCurrentAssistantMessage((prev) => prev + data.chunk);
+      setCurrentAssistantMessage((prev) => prev + message.payload.chunk);
    });
    // Listen for transcription results
-    const unsubTranscription = ws.on(
+    const unsubTranscription = ws.on("transcription_result", () => {
-      "transcription_result",
+      // Note: Transcription is already broadcast as activity_log with type "transcript"
-      (_payload: unknown) => {
+      // No need to log it again here to avoid duplication
-        // Note: Transcription is already broadcast as activity_log with type "transcript"
+      setIsProcessingAudio(false);
-        // No need to log it again here to avoid duplication
+    });
        setIsProcessingAudio(false);
      }
    );
    // Listen for artifacts
-    const unsubArtifact = ws.on("artifact", (payload: unknown) => {
+    const unsubArtifact = ws.on("artifact", (message) => {
-      const artifact = payload as Artifact;
+      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 unsubAudioOutput = ws.on("audio_output", async (message) => {
-      const data = payload as { audio: string; format: string; id: string };
+      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",
+          type: "session",
-          payload: { id: data.id },
+          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",
+          type: "session",
-          payload: {
+          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",
+                type: "session",
-                payload: {},
+                message: {
                  type: "abort_request",
                },
              });
            },
            onSpeechEnd: async (audioData: Float32Array) => {
@@ -462,8 +458,9 @@ function App() {
                );
                ws.send({
-                  type: "audio_chunk",
+                  type: "session",
-                  payload: {
+                  message: {
                    type: "audio_chunk",
                    audio: base64Audio,
                    format: audioBlob.type,
                    isLast: true,
@@ -1,16 +1,14 @@
 import { useEffect, useRef, useState, useCallback } from 'react';
-
+import type {
-interface WebSocketMessage {
+  WSInboundMessage,
-  type: string;
+  WSOutboundMessage,
-  payload: unknown;
+  SessionOutboundMessage
-}
+} from '../../server/messages.js';
 type MessageHandler = (payload: unknown) => void;
 export interface UseWebSocketReturn {
  isConnected: boolean;
-  send: (message: WebSocketMessage) => void;
+  send: (message: WSInboundMessage) => void;
-  on: (type: string, handler: MessageHandler) => () => 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<WebSocket | null>(null);
-  const handlersRef = useRef<Map<string, Set<MessageHandler>>>(new Map());
+  const handlersRef = useRef<Map<SessionOutboundMessage['type'], Set<(message: SessionOutboundMessage) => void>>>(new Map());
  const reconnectTimeoutRef = useRef<number>();
  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;
+          const wsMessage: WSOutboundMessage = JSON.parse(event.data);
          console.log(`[WS] Received message type: ${message.type}`, message.payload);
-          // Call all registered handlers for this message type
+          // Only session messages trigger handlers
-          const handlers = handlersRef.current.get(message.type);
+          if (wsMessage.type === 'session') {
-          if (handlers) {
+            const sessionMessage = wsMessage.message;
-            handlers.forEach((handler) => {
+            console.log(`[WS] Received session message type: ${sessionMessage.type}`);
-              try {
+
-                handler(message.payload);
+            // Call all registered handlers for this message type
-              } catch (err) {
+            const handlers = handlersRef.current.get(sessionMessage.type);
-                console.error(`[WS] Error in handler for ${message.type}:`, err);
+            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]
  );