#Build a voice assistant and text to speech

// protocol.ts: imported by the browser and the server. Audio is 24 kHz mono PCM16.
export type VoiceIn =
  // The user started talking: stop the answer.
  | { type: 'speech-started'; itemId: string }
  // A chunk of the answer's audio, as base64.
  | { type: 'audio'; itemId: string; audio: string }
  // What the user said, once it's transcribed.
  | { type: 'heard'; itemId: string; text: string }
  // The next words of the answer's transcript.
  | { type: 'said'; itemId: string; delta: string }
  // Text to speech: a chunk of audio, the end, or a failure.
  | { type: 'speech'; id: number; audio: string }
  | { type: 'speech-end'; id: number }
  | { type: 'speech-failed'; id: number; error: string };

export type VoiceOut =
  // A microphone frame, as binary.
  | ArrayBuffer
  // Stop answering. `heardMs` is how much of `itemId` the user heard.
  | { type: 'interrupt'; itemId?: string; heardMs?: number }
  | { type: 'speak'; id: number; text: string; voice: string }
  | { type: 'stop-speaking'; id: number };

// mic.ts
import { voice } from './live';

/** Start sending the microphone, and return a function that stops it. */
export async function startMic(): Promise<() => void> {
  const stream = await navigator.mediaDevices.getUserMedia({
    // Without echo cancellation, the answer from the speakers reaches the microphone, and the
    // server hears the user interrupting.
    audio: { echoCancellation: true, noiseSuppression: true, channelCount: 1 }
  });
  // At the device's rate: some browsers can't connect a microphone to a context at another one.
  const context = new AudioContext();
  await context.audioWorklet.addModule('/mic-worklet.js');
  const source = context.createMediaStreamSource(stream);
  const mic = new AudioWorkletNode(context, 'mic', { numberOfOutputs: 0 });
  mic.port.onmessage = ({ data }: MessageEvent<ArrayBuffer>) => {
    // Dropped while disconnected, and while 64 KB are already waiting to go out.
    voice.send(data, { volatile: true, maxBuffered: 64 * 1024 });
  };
  source.connect(mic);

  return () => {
    source.disconnect();
    mic.port.onmessage = null;
    for (const track of stream.getTracks()) track.stop();
    void context.close();
  };
}

// live.ts
import {
  createWebSocket,
  createWebSocketRequests,
  createWebSocketState,
  retryableClose
} from 'react-websocket';
import { playback } from './playback';
import type { VoiceIn, VoiceOut } from './protocol';

// The conversation's id, for this tab: a reconnect or a reload within 30 s picks it up again.
const session = sessionStorage.getItem('voice-session') ?? crypto.randomUUID();
sessionStorage.setItem('voice-session', session);

const scheme = location.protocol === 'https:' ? 'wss:' : 'ws:';

export const voice = createWebSocket<VoiceIn, VoiceOut>(
  () => `${scheme}//${location.host}/voice?session=${session}`,
  {
    binaryType: 'arraybuffer',
    // Microphone frames go out as binary, everything else as JSON. The server sends only JSON.
    serialize: message => (message instanceof ArrayBuffer ? message : JSON.stringify(message)),
    parse: data => JSON.parse(data as string),
    // 4409: the conversation moved to another tab. Don't take it back.
    reconnect: { shouldReconnect: close => close.code !== 4409 && retryableClose(close) }
  }
);

export type Line = { itemId: string; role: 'user' | 'assistant'; text: string; cut?: boolean };

export const transcript = createWebSocketState(voice, {
  initialState: [] as Line[],
  reduce: (lines, message): Line[] => {
    switch (message.type) {
      case 'speech-started':
        return [...lines, { itemId: message.itemId, role: 'user', text: '' }];
      case 'heard':
        return lines.map(line =>
          line.itemId === message.itemId ? { ...line, text: message.text } : line
        );
      case 'said':
        return lines.some(line => line.itemId === message.itemId)
          ? lines.map(line =>
              line.itemId === message.itemId ? { ...line, text: line.text + message.delta } : line
            )
          : [...lines, { itemId: message.itemId, role: 'assistant', text: message.delta }];
      default:
        return lines;
    }
  },
  batch: 'frame'
});

/** Stop the answer: what's playing, what the model is still generating, and its transcript. */
export function interrupt() {
  const heard = playback.stop();
  voice.send({ type: 'interrupt', ...heard }, { priority: true });
  if (!heard) return;
  transcript.setState(lines =>
    lines.map(line => (line.itemId === heard.itemId ? { ...line, cut: true } : line))
  );
}

voice.on('message', message => {
  if (message.type === 'audio') playback.play(message.itemId, message.audio);
  // The server heard the user start talking over the answer.
  if (message.type === 'speech-started') interrupt();
});

// The server gives up on an answer whose connection drops: stop playing it too.
voice.on('close', () => playback.stop());

export const speech = createWebSocketRequests(voice, {
  send: ({ text, voice }: { text: string; voice: string }, id) => ({
    type: 'speak',
    id,
    text,
    voice
  }),
  reply: message =>
    message.type === 'speech' || message.type === 'speech-end' || message.type === 'speech-failed'
      ? message.id
      : undefined,
  error: message => (message.type === 'speech-failed' ? new Error(message.error) : undefined),
  done: message => message.type === 'speech-end',
  cancel: id => ({ type: 'stop-speaking', id }),
  timeout: 30_000
});

// wav.ts: 24 kHz mono PCM16 chunks, as a WAV file.
export function toWav(chunks: string[]): Blob {
  const parts = chunks.map(chunk => Uint8Array.from(atob(chunk), char => char.charCodeAt(0)));
  const size = parts.reduce((total, part) => total + part.length, 0);
  const header = new DataView(new ArrayBuffer(44));
  const write = (offset: number, text: string) => {
    for (let i = 0; i < text.length; i++) header.setUint8(offset + i, text.charCodeAt(i));
  };
  write(0, 'RIFF');
  header.setUint32(4, 36 + size, true);
  write(8, 'WAVEfmt ');
  header.setUint32(16, 16, true); // the format chunk's size
  header.setUint16(20, 1, true); // PCM
  header.setUint16(22, 1, true); // mono
  header.setUint32(24, 24_000, true); // samples per second
  header.setUint32(28, 48_000, true); // bytes per second
  header.setUint16(32, 2, true); // bytes per sample
  header.setUint16(34, 16, true); // bits per sample
  write(36, 'data');
  header.setUint32(40, size, true);
  return new Blob([header, ...parts], { type: 'audio/wav' });
}

// server.ts
import { createServer } from 'node:http';
import { WebSocket, WebSocketServer } from 'ws';
import { userFromCookie } from './auth.ts';
import type { VoiceIn, VoiceOut } from './protocol.ts';

// Check OpenAI's model list for the current realtime model.
const REALTIME_URL = 'wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1';
const headers = { Authorization: `Bearer ${process.env.OPENAI_API_KEY}` };

/** The Realtime API's events this server reads. */
type RealtimeEvent =
  | { type: 'response.created' | 'response.done' }
  | { type: 'input_audio_buffer.speech_started'; item_id: string }
  | {
      type: 'conversation.item.input_audio_transcription.completed';
      item_id: string;
      transcript: string;
    }
  | { type: 'response.output_audio.delta'; item_id: string; delta: string }
  | { type: 'response.output_audio_transcript.delta'; item_id: string; delta: string }
  | { type: 'error'; error: { message: string } };

type Session = {
  userId: string;
  upstream: WebSocket; // to OpenAI
  socket?: WebSocket; // to the browser, while it's connected
  responding: boolean;
  sentMs: Map<string, number>; // how much audio of each item went to the browser
  idle?: ReturnType<typeof setTimeout>;
};

const sessions = new Map<string, Session>();

function send(session: Session, message: VoiceIn) {
  session.socket?.send(JSON.stringify(message));
}

function openSession(id: string, userId: string): Session {
  const upstream = new WebSocket(REALTIME_URL, { headers });
  const session: Session = { userId, upstream, responding: false, sentMs: new Map() };
  sessions.set(id, session);

  upstream.on('open', () => {
    upstream.send(
      JSON.stringify({
        type: 'session.update',
        session: {
          type: 'realtime',
          instructions: 'You are a friendly voice assistant. Keep your answers short.',
          audio: {
            input: {
              format: { type: 'audio/pcm', rate: 24_000 },
              transcription: { model: 'gpt-4o-mini-transcribe' },
              // Answers when the user stops talking, and stops answering when they start.
              turn_detection: { type: 'semantic_vad' }
            },
            output: { format: { type: 'audio/pcm', rate: 24_000 }, voice: 'marin' }
          }
        }
      })
    );
  });

  upstream.on('message', data => {
    const event = JSON.parse(String(data)) as RealtimeEvent;
    switch (event.type) {
      case 'response.created':
      case 'response.done':
        session.responding = event.type === 'response.created';
        break;
      case 'input_audio_buffer.speech_started':
        send(session, { type: 'speech-started', itemId: event.item_id });
        break;
      case 'conversation.item.input_audio_transcription.completed':
        send(session, { type: 'heard', itemId: event.item_id, text: event.transcript });
        break;
      case 'response.output_audio.delta': {
        const ms = Buffer.byteLength(event.delta, 'base64') / 48; // 24 kHz, 2 bytes a sample
        session.sentMs.set(event.item_id, (session.sentMs.get(event.item_id) ?? 0) + ms);
        send(session, { type: 'audio', itemId: event.item_id, audio: event.delta });
        break;
      }
      case 'response.output_audio_transcript.delta':
        send(session, { type: 'said', itemId: event.item_id, delta: event.delta });
        break;
      case 'error':
        console.error(event.error.message);
        break;
    }
  });

  upstream.on('close', () => {
    sessions.delete(id);
    session.socket?.close(1011, 'Session ended'); // the browser reconnects to a new one
  });
  return session;
}

type Interrupt = Extract<VoiceOut, { type: 'interrupt' }>;
type Speak = Extract<VoiceOut, { type: 'speak' }>;

function interrupt(session: Session, { itemId, heardMs = 0 }: Interrupt) {
  const upstream = session.upstream;
  if (upstream.readyState !== WebSocket.OPEN) return;
  if (session.responding) upstream.send(JSON.stringify({ type: 'response.cancel' }));
  // Cut the answer back to what the user heard, so the model doesn't think they heard the rest.
  if (itemId && heardMs < (session.sentMs.get(itemId) ?? 0)) {
    const truncate = { type: 'conversation.item.truncate', item_id: itemId, content_index: 0 };
    upstream.send(JSON.stringify({ ...truncate, audio_end_ms: Math.floor(heardMs) }));
  }
}

/** Text to speech, streamed: raw 24 kHz PCM16 as it's generated. */
async function speak(session: Session, { id, text, voice }: Speak, signal: AbortSignal) {
  const response = await fetch('https://api.openai.com/v1/audio/speech', {
    method: 'POST',
    headers: { ...headers, 'Content-Type': 'application/json' },
    body: JSON.stringify({ model: 'gpt-4o-mini-tts', voice, input: text, response_format: 'pcm' }),
    signal
  });
  if (!response.ok || !response.body) throw new Error(`HTTP ${response.status}`);
  let odd = Buffer.alloc(0); // a byte of a sample split across two chunks
  for await (const chunk of response.body) {
    const bytes = Buffer.concat([odd, chunk]);
    const whole = bytes.length - (bytes.length % 2);
    odd = bytes.subarray(whole);
    if (!whole) continue;
    send(session, { type: 'speech', id, audio: bytes.subarray(0, whole).toString('base64') });
  }
  send(session, { type: 'speech-end', id });
}

function connected(socket: WebSocket, session: Session) {
  clearTimeout(session.idle);
  // The conversation opened somewhere else, such as a duplicated tab: it moves there.
  session.socket?.close(4409, 'Opened in another tab');
  session.socket = socket;
  const speaking = new Map<number, AbortController>();

  socket.on('message', (data, isBinary) => {
    const { upstream } = session;
    if (isBinary) {
      // A microphone frame. Dropped until OpenAI's socket is open, like the browser drops them.
      if (upstream.readyState !== WebSocket.OPEN) return;
      const audio = Buffer.from(data as Buffer).toString('base64');
      return upstream.send(JSON.stringify({ type: 'input_audio_buffer.append', audio }));
    }
    const message = JSON.parse(String(data)) as Exclude<VoiceOut, ArrayBuffer>;
    switch (message.type) {
      case 'interrupt':
        interrupt(session, message);
        break;
      case 'speak': {
        const controller = new AbortController();
        speaking.set(message.id, controller);
        speak(session, message, controller.signal)
          .catch(error => {
            if (controller.signal.aborted) return;
            console.error(error);
            send(session, { type: 'speech-failed', id: message.id, error: 'Speech failed.' });
          })
          .finally(() => speaking.delete(message.id));
        break;
      }
      case 'stop-speaking':
        speaking.get(message.id)?.abort();
        break;
    }
  });

  socket.on('close', () => {
    for (const controller of speaking.values()) controller.abort();
    if (session.socket !== socket) return;
    session.socket = undefined;
    // Keep the conversation for 30 s, for a reconnect or a reload.
    session.idle = setTimeout(() => session.upstream.close(), 30_000);
  });
}

const server = createServer();
const wss = new WebSocketServer({ noServer: true });

server.on('upgrade', async (request, stream, head) => {
  const url = new URL(request.url ?? '/', 'http://localhost');
  if (url.pathname !== '/voice' || request.headers.origin !== process.env.APP_ORIGIN) {
    return stream.destroy();
  }
  const userId = await userFromCookie(request.headers.cookie);
  const id = url.searchParams.get('session') ?? '';
  wss.handleUpgrade(request, stream, head, socket => {
    if (!userId || !id) return socket.close(1008, 'Unauthorized');
    const existing = sessions.get(id);
    if (existing && existing.userId !== userId) return socket.close(1008, 'Unauthorized');
    connected(socket, existing ?? openSession(id, userId));
  });
});

server.listen(Number(process.env.PORT ?? 8080));