Blame for server/voice-ws.ts

/**
 * Persistent real-time Arabic speech bridge.
 *
 * The mobile/web client streams sequenced PCM16 frames here. One persistent
 * recognizer WebSocket is retained for each logical session, including across
 * short client reconnects. No live audio is sent through Whisper HTTP.
 */
import { WebSocket, WebSocketServer } from "ws";
import type { Server } from "http";
import {
  OrderedAudioFrames,
  type BufferedAudioFrame,
} from "../shared/realtime-voice-protocol";
import {
  resolveSttProvider,
  type RealtimeSttRecognizer,
  type SendVoiceMessage,
} from "./realtime-stt-provider";
import { ShihabRealtimeRecognizer } from "./shihab-realtime-recognizer";

const SESSION_RETENTION_MS = 30000;
const UPSTREAM_RECONNECT_INITIAL_MS = 250;
const UPSTREAM_RECONNECT_MAX_MS = 8000;
const MAX_UPSTREAM_REPLAY_FRAMES = 600;

type SendMessage = SendVoiceMessage;

interface PcmStartMessage {
  type: "pcm_start";
  sessionId: string;
  sampleRate?: number;
  firstAudioSequence?: number;
}

interface PcmMessage {
  type: "pcm";
  sessionId: string;
  audioSequence: number;
  audio: string;
}

interface PcmStopMessage {
  type: "pcm_stop";
  sessionId: string;
  lastAudioSequence?: number;
}

function sendJson(socket: WebSocket, message: Record<string, unknown>): void {
  if (socket.readyState === WebSocket.OPEN) {
    socket.send(JSON.stringify(message));
  }
}

function recognizerWebSocketUrl(): string {
  const configured =
    process.env.WHISPER_VAD_WS_URL ??
    process.env.WHISPER_VAD_URL ??
    process.env.EXPO_PUBLIC_WHISPER_VAD_URL ??
    "https://taqarub-vosk-ar.hf.space";
  const raw = configured
    .trim()
    .replace(/^['"]|['"]$/g, "")
    .replace(/\/+$/, "");
  if (!raw) {
    throw new Error(
      "Persistent recognizer is not configured. Set WHISPER_VAD_URL.",
    );
  }
  if (raw.endsWith("/ws/transcribe")) {
    return raw.replace(/^https:/, "wss:").replace(/^http:/, "ws:");
  }
  return `${raw.replace(/^https:/, "wss:").replace(/^http:/, "ws:")}/ws/transcribe`;
}

function transcriptWords(text: string): string[] {
  return text
    .replace(/[^\p{L}\p{N}]+/gu, " ")
    .trim()
    .split(/\s+/)
    .filter(Boolean);
}

class PersistentArabicRecognizer implements RealtimeSttRecognizer {
  private socket: WebSocket | null = null;
  private connectPromise: Promise<void> | null = null;
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  private reconnectAttempt = 0;
  private handshakeReady = false;
  private active = true;
  private lastAcknowledgedSequence = -1;
  private replayFrames = new Map<number, BufferedAudioFrame>();
  private latestRevision = new Map<string, number>();
  private finalizedUtterances = new Set<string>();
  private finalEventHistory = new Map<number, Record<string, unknown>>();
  private latestInterimEvent: Record<string, unknown> | null = null;
  private nextUtteranceIndex = 0;
  private upstreamProtocolVersion = 1;
  private legacyRevision = 0;
  private legacyInterim = "";

  constructor(
    readonly sessionId: string,
    private readonly sampleRate: number,
    private readonly emit: SendMessage,
  ) {}

  start(): Promise<void> {
    this.active = true;
    return this.ensureConnected();
  }

  accept(frame: BufferedAudioFrame): void {
    if (frame.audioSequence <= this.lastAcknowledgedSequence) {
      this.emitAck();
      return;
    }
    if (!this.replayFrames.has(frame.audioSequence)) {
      if (this.replayFrames.size >= MAX_UPSTREAM_REPLAY_FRAMES) {
        this.emit({
          type: "error",
          sessionId: this.sessionId,
          message: "Recognizer replay buffer is full",
        });
        return;
      }
      this.replayFrames.set(frame.audioSequence, frame);
    }
    this.sendFrame(frame);
    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
      void this.ensureConnected().catch(() => this.scheduleReconnect());
    }
  }

  stop(): void {
    this.active = false;
    if (this.reconnectTimer) clearTimeout(this.reconnectTimer);
    this.reconnectTimer = null;
    if (this.socket?.readyState === WebSocket.OPEN) {
      this.socket.send(
        JSON.stringify({
          type: "pcm_stop",
          sessionId: this.sessionId,
        }),
      );
    }
  }

  close(): void {
    this.active = false;
    if (this.reconnectTimer) clearTimeout(this.reconnectTimer);
    this.reconnectTimer = null;
    this.socket?.close();
    this.socket = null;
    this.connectPromise = null;
  }

  acknowledgedSequence(): number {
    return this.lastAcknowledgedSequence;
  }

  replayTranscriptEvents(): void {
    for (const event of [...this.finalEventHistory.entries()]
      .sort(([left], [right]) => left - right)
      .map(([, value]) => value)) {
      this.emit(event);
    }
    if (this.latestInterimEvent) this.emit(this.latestInterimEvent);
  }

  private ensureConnected(): Promise<void> {
    if (this.socket?.readyState === WebSocket.OPEN && this.handshakeReady) {
      return Promise.resolve();
    }
    if (this.connectPromise) return this.connectPromise;

    const promise = new Promise<void>((resolve, reject) => {
      let url: string;
      try {
        url = recognizerWebSocketUrl();
      } catch (error) {
        reject(error);
        return;
      }

      const socket = new WebSocket(url, { perMessageDeflate: false });
      this.socket = socket;
      this.handshakeReady = false;
      let connected = false;
      const timeout = setTimeout(() => {
        if (connected) return;
        socket.close();
        reject(
          new Error(`Timed out connecting to Arabic recognizer at ${url}`),
        );
      }, 8000);

      socket.on("message", (raw) => {
        let message: Record<string, any>;
        try {
          message = JSON.parse(raw.toString());
        } catch {
          return;
        }

        if (message.type === "ready") {
          this.upstreamProtocolVersion = Number(message.protocolVersion ?? 1);
          const bufferedSequences = [...this.replayFrames.keys()];
          const firstAudioSequence =
            bufferedSequences.length > 0
              ? Math.min(...bufferedSequences)
              : this.lastAcknowledgedSequence + 1;
          socket.send(
            JSON.stringify({
              type: "pcm_start",
              sessionId: this.sessionId,
              sampleRate: this.sampleRate,
              firstAudioSequence,
              firstUtteranceIndex: this.nextUtteranceIndex,
              language: "ar",
            }),
          );
          return;
        }

        if (message.type === "pcm_ready") {
          connected = true;
          clearTimeout(timeout);
          this.handshakeReady = true;
          this.reconnectAttempt = 0;
          this.replay();
          resolve();
          return;
        }

        if (message.type === "audio_ack") {
          const sequence = Number(message.audioSequence);
          if (Number.isFinite(sequence)) {
            this.lastAcknowledgedSequence = Math.max(
              this.lastAcknowledgedSequence,
              sequence,
            );
            for (const key of this.replayFrames.keys()) {
              if (key <= this.lastAcknowledgedSequence) {
                this.replayFrames.delete(key);
              }
            }
            this.emitAck();
          }
          return;
        }

        if (message.type === "interim" || message.type === "final") {
          this.forwardTranscript(message);
          return;
        }

        if (
          message.type === "partial_transcript" ||
          message.type === "final_transcript"
        ) {
          this.forwardLegacyTranscript(message);
          return;
        }

        if (message.type === "error") {
          this.emit({
            type: "error",
            sessionId: this.sessionId,
            message: String(message.message ?? "Recognizer error"),
          });
        }
      });

      socket.on("error", (error) => {
        if (!connected) {
          clearTimeout(timeout);
          reject(error);
        }
      });

      socket.on("close", () => {
        clearTimeout(timeout);
        if (this.socket === socket) this.socket = null;
        this.handshakeReady = false;
        this.connectPromise = null;
        if (this.active) this.scheduleReconnect();
      });
    }).finally(() => {
      this.connectPromise = null;
    });

    this.connectPromise = promise;
    return promise;
  }

  private scheduleReconnect(): void {
    if (!this.active || this.reconnectTimer) return;
    const delay = Math.min(
      UPSTREAM_RECONNECT_MAX_MS,
      UPSTREAM_RECONNECT_INITIAL_MS * 2 ** this.reconnectAttempt++,
    );
    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      void this.ensureConnected().catch(() => this.scheduleReconnect());
    }, delay);
  }

  private sendFrame(frame: BufferedAudioFrame): void {
    if (!this.handshakeReady || this.socket?.readyState !== WebSocket.OPEN) {
      return;
    }
    this.socket.send(
      JSON.stringify({
        type: "pcm",
        sessionId: this.sessionId,
        audioSequence: frame.audioSequence,
        chunkId: frame.audioSequence,
        sampleRate: this.sampleRate,
        audio: frame.audio,
      }),
    );
    if (this.upstreamProtocolVersion < 2) {
      this.lastAcknowledgedSequence = Math.max(
        this.lastAcknowledgedSequence,
        frame.audioSequence,
      );
      this.replayFrames.delete(frame.audioSequence);
      this.emitAck();
    }
  }

  private replay(): void {
    const frames = [...this.replayFrames.values()].sort(
      (left, right) => left.audioSequence - right.audioSequence,
    );
    for (const frame of frames) this.sendFrame(frame);
  }

  private emitAck(): void {
    this.emit({
      type: "audio_ack",
      sessionId: this.sessionId,
      audioSequence: this.lastAcknowledgedSequence,
    });
  }

  private forwardTranscript(message: Record<string, any>): void {
    const utteranceId = String(message.utteranceId ?? "");
    const revision = Number(message.revision);
    const utteranceIndex = Number(message.utteranceIndex);
    const text = String(message.text ?? "").trim();
    if (
      !utteranceId ||
      !Number.isFinite(revision) ||
      !Number.isFinite(utteranceIndex)
    ) {
      return;
    }

    const previousRevision = this.latestRevision.get(utteranceId) ?? -1;
    if (message.type === "interim") {
      if (
        this.finalizedUtterances.has(utteranceId) ||
        revision <= previousRevision
      ) {
        return;
      }
    } else {
      if (
        this.finalizedUtterances.has(utteranceId) ||
        revision < previousRevision
      ) {
        return;
      }
      this.finalizedUtterances.add(utteranceId);
    }

    this.latestRevision.set(utteranceId, revision);
    const event: Record<string, unknown> = {
      type: message.type,
      sessionId: this.sessionId,
      utteranceId,
      utteranceIndex,
      revision,
      audioSequence: Number(message.audioSequence),
      text,
      words: message.type === "final" ? transcriptWords(text) : undefined,
    };
    if (message.type === "final") {
      this.nextUtteranceIndex = Math.max(
        this.nextUtteranceIndex,
        utteranceIndex + 1,
      );
      this.finalEventHistory.set(utteranceIndex, event);
      if (this.latestInterimEvent?.utteranceId === utteranceId) {
        this.latestInterimEvent = null;
      }
    } else {
      this.latestInterimEvent = event;
    }
    this.emit(event);
  }

  private forwardLegacyTranscript(message: Record<string, any>): void {
    const text = String(message.text ?? "").trim();
    if (!text) return;

    if (message.type === "partial_transcript") {
      if (text === this.legacyInterim) return;
      this.legacyInterim = text;
      this.legacyRevision++;
      this.forwardTranscript({
        type: "interim",
        utteranceId: `${this.sessionId}:${this.nextUtteranceIndex}`,
        utteranceIndex: this.nextUtteranceIndex,
        revision: this.legacyRevision,
        audioSequence: Number(message.chunkId ?? -1),
        text,
      });
      return;
    }

    this.legacyRevision++;
    this.forwardTranscript({
      type: "final",
      utteranceId: `${this.sessionId}:${this.nextUtteranceIndex}`,
      utteranceIndex: this.nextUtteranceIndex,
      revision: this.legacyRevision,
      audioSequence: Number(message.chunkId ?? -1),
      text,
    });
    this.legacyRevision = 0;
    this.legacyInterim = "";
  }
}

function createRecognizer(
  sessionId: string,
  sampleRate: number,
  emit: SendMessage,
): RealtimeSttRecognizer {
  const provider = resolveSttProvider();
  if (provider === "shihab") {
    return new ShihabRealtimeRecognizer(sessionId, emit);
  }
  return new PersistentArabicRecognizer(sessionId, sampleRate, emit);
}

class LiveRecognitionSession {
  private sender: SendMessage = () => {};
  private orderedAudio: OrderedAudioFrames;
  private recognizer: RealtimeSttRecognizer;
  private retentionTimer: ReturnType<typeof setTimeout> | null = null;

  constructor(
    readonly sessionId: string,
    sampleRate: number,
    firstAudioSequence: number,
    private readonly onExpire: () => void,
  ) {
    this.orderedAudio = new OrderedAudioFrames(firstAudioSequence);
    this.recognizer = createRecognizer(sessionId, sampleRate, (message) =>
      this.sender(message),
    );
  }

  async bind(sender: SendMessage): Promise<void> {
    if (this.retentionTimer) clearTimeout(this.retentionTimer);
    this.retentionTimer = null;
    this.sender = sender;
    await this.recognizer.start();
    this.sender({
      type: "pcm_ready",
      sessionId: this.sessionId,
      acknowledgedAudioSequence: this.recognizer.acknowledgedSequence(),
    });
    this.recognizer.replayTranscriptEvents();
  }

  receive(message: PcmMessage): void {
    const result = this.orderedAudio.receive({
      audioSequence: message.audioSequence,
      audio: message.audio,
    });
    if (result.duplicate) {
      this.sender({
        type: "audio_ack",
        sessionId: this.sessionId,
        audioSequence: this.recognizer.acknowledgedSequence(),
      });
      return;
    }
    for (const frame of result.ready) this.recognizer.accept(frame);
  }

  stop(): void {
    this.recognizer.stop();
  }

  unbind(sender: SendMessage): void {
    if (this.sender !== sender) return;
    this.sender = () => {};
    if (this.retentionTimer) clearTimeout(this.retentionTimer);
    this.retentionTimer = setTimeout(() => {
      this.recognizer.close();
      this.onExpire();
    }, SESSION_RETENTION_MS);
  }

  closeSoon(): void {
    if (this.retentionTimer) clearTimeout(this.retentionTimer);
    this.retentionTimer = setTimeout(() => {
      this.recognizer.close();
      this.onExpire();
    }, 5000);
  }
}

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

function getOrCreateSession(message: PcmStartMessage): LiveRecognitionSession {
  const existing = sessions.get(message.sessionId);
  if (existing) return existing;

  const session = new LiveRecognitionSession(
    message.sessionId,
    message.sampleRate ?? 16000,
    message.firstAudioSequence ?? 0,
    () => sessions.delete(message.sessionId),
  );
  sessions.set(message.sessionId, session);
  return session;
}

export function registerVoiceWebSocket(httpServer: Server): void {
  const provider = resolveSttProvider();
  if (
    (process.env.STT_PROVIDER ?? "shihab").trim().toLowerCase() === "shihab" &&
    provider === "existing"
  ) {
    console.warn(
      "[voice-ws] SHIHAB_API_KEY is missing; using the existing STT provider",
    );
  }
  const server = new WebSocketServer({
    server: httpServer,
    path: "/ws/voice",
    perMessageDeflate: false,
  });

  server.on("connection", (socket) => {
    const boundSessionIds = new Set<string>();
    const send: SendMessage = (message) => sendJson(socket, message);
    send({ type: "ready", protocolVersion: 2 });

    socket.on("message", (raw) => {
      let message: Record<string, any>;
      try {
        message = JSON.parse(raw.toString());
      } catch {
        return;
      }

      if (message.type === "ping") {
        send({ type: "pong" });
        return;
      }

      if (message.type === "pcm_start") {
        const start = message as PcmStartMessage;
        if (!start.sessionId) {
          send({ type: "error", message: "sessionId is required" });
          return;
        }
        const session = getOrCreateSession(start);
        boundSessionIds.add(start.sessionId);
        void session.bind(send).catch((error) => {
          send({
            type: "error",
            sessionId: start.sessionId,
            message: error instanceof Error ? error.message : String(error),
          });
        });
        return;
      }

      if (message.type === "pcm") {
        const pcm = message as PcmMessage;
        const session = sessions.get(pcm.sessionId);
        if (!session) {
          send({
            type: "error",
            sessionId: pcm.sessionId,
            message: "Unknown PCM session; send pcm_start first",
          });
          return;
        }
        session.receive(pcm);
        return;
      }

      if (message.type === "pcm_stop") {
        const stop = message as PcmStopMessage;
        const session = sessions.get(stop.sessionId);
        session?.stop();
        session?.closeSoon();
        return;
      }

      if (message.type === "chunk") {
        send({
          type: "error",
          message:
            "Recorded audio chunks are no longer supported; stream PCM continuously.",
        });
      }
    });

    socket.on("close", () => {
      for (const sessionId of boundSessionIds) {
        sessions.get(sessionId)?.unbind(send);
      }
    });
  });

  console.log(
    `[voice-ws] persistent PCM WebSocket registered at /ws/voice (provider=${provider})`,
  );
}