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.
This commit is contained in:
Mohamed Boudra
2025-10-20 15:28:28 +02:00
parent 8365fb889b
commit af67235997
10 changed files with 1114 additions and 937 deletions

View File

@@ -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<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,
})),
};
}

View File

@@ -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
}
}

View File

@@ -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<string, PendingPlayback>();
export class TTSManager {
private pendingPlaybacks: Map<string, PendingPlayback> = 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<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;
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<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);
}
}
}

View File

@@ -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}`

View File

@@ -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,
};
}

View File

@@ -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();
}
}

View File

@@ -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)
}

View File

@@ -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<WebSocket, ClientSession> = new Map();
private conversationIdToWs: Map<string, WebSocket> = new Map(); // Reverse map for phase management
private messageHandler?: (conversationId: string, message: string, abortSignal: AbortSignal) => Promise<void>;
private audioHandler?: (conversationId: string, audio: Buffer, format: string, abortSignal: AbortSignal) => Promise<string>;
private sessions: Map<WebSocket, Session> = new Map();
private conversationIdToWs: Map<string, WebSocket> = 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<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: {} });
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<void> {
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<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
}
/**
* Close the WebSocket server
*/
public close(): void {
this.sessions.forEach((_session, ws) => {
this.sessions.forEach((session, ws) => {
session.cleanup();
ws.close();
});
this.wss.close();

View File

@@ -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<string, unknown>;
};
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,

View File

@@ -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<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;
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]
);