s2s-client.js

// @ts-check
/**
 * Minimal WebRTC client for the Hugging Face speech-to-speech load balancer.
 *
 * Two-step handshake:
 *
 *   1. POST `<lb_url>/session` -> JSON `{ connect_url: wss://<compute>/v1/realtime?session_token=<JWT>, ... }`
 *   2. POST that connect_url (after rewriting `wss://` -> `https://`) with
 *      `Content-Type: application/sdp` and the SDP offer body. The backend
 *      (aiortc) answers `201 application/sdp` with the SDP answer.
 *
 * Once the peer connection is up:
 *
 *   - Audio is bidirectional on the media tracks (the browser captures the
 *     user's mic and sends Opus 48 kHz; the backend resamples both ways
 *     internally to feed the 16 kHz STT pipeline and the 16 kHz TTS output).
 *   - JSON events flow on a data channel labelled `oai-events`, using the
 *     OpenAI Realtime protocol (session.update, response.audio_transcript.delta,
 *     conversation.item.input_audio_transcription.delta, ...).
 *
 * References:
 *   https://github.com/huggingface/speech-to-speech (feat/webrtc-transport branch)
 *   https://platform.openai.com/docs/guides/realtime-webrtc
 *
 * @typedef {"idle" | "creating-session" | "connecting" | "connected" |
 *           "user-speaking" | "processing" | "ai-speaking" | "closed" | "error"
 * } S2sStatus
 *
 * @typedef {Object} S2sSessionInfo
 * @property {string} sessionId
 * @property {string} connectUrl
 * @property {string} websocketUrl
 * @property {string} sessionToken
 * @property {number} pendingTimeoutS
 *
 * @typedef {Object} S2sOptions
 * @property {string} loadBalancerUrl  Base URL of the LB (e.g. `https://...endpoints.huggingface.cloud`)
 * @property {string} voice            Qwen3-TTS speaker name (Aiden, Ryan, Dylan, Eric, Ono_Anna, Serena, Sohee, Uncle_Fu, Vivian)
 * @property {string} instructions     System prompt sent in `session.update`
 * @property {MediaStreamTrack} inputTrack
 *
 * @typedef {Object} TranscriptEvent
 * @property {"user" | "assistant"} role
 * @property {string} text
 * @property {boolean} partial
 */

const DEFAULT_PLAYOUT_DELAY_S = 0.15;

// Opus fmtp knobs forced via SDP munging. Inline FEC trades a few bytes for
// loss resilience; DTX off keeps comfort-noise frames flowing so the server's
// VAD never thinks we silently dropped.
const OPUS_FMTP_OVERRIDES = {
  minptime: "10",
  useinbandfec: "1",
  usedtx: "0",
  stereo: "0",
  "sprop-stereo": "0",
};

const ICE_DISCONNECT_GRACE_MS = 5_000;

export class S2sRealtimeClient extends EventTarget {
  /** @param {S2sOptions} options */
  constructor(options) {
    super();
    /** @type {S2sOptions} */
    this.options = options;
    /** @type {RTCPeerConnection | null} */
    this._pc = null;
    /** @type {RTCDataChannel | null} */
    this._dc = null;
    /** @type {S2sStatus} */
    this._status = "idle";
    /** @type {number | null} */
    this._iceGraceTimer = null;
    /** @type {(() => void) | null} */
    this._pendingVisibilityHandler = null;
  }

  get status() {
    return this._status;
  }

  /** @param {S2sStatus} status */
  _setStatus(status) {
    if (this._status === status) return;
    this._status = status;
    this.dispatchEvent(new CustomEvent("status", { detail: { status } }));
  }

  _markAudible() {
    if (this._status === "ai-speaking") return;
    if (this._status === "closed" || this._status === "error") return;
    this._setStatus("ai-speaking");
  }

  /**
   * Full handshake:
   *   1. Create a session on the load balancer
   *   2. Open the WebRTC peer connection against the routed compute endpoint
   *   3. Wait for the SDP answer + setRemoteDescription
   *
   * Resolves once the SDP exchange completes; the data channel and the
   * `session.update` payload are sent asynchronously when the channel opens.
   *
   * @returns {Promise<void>}
   */
  async connect() {
    if (this._pc) throw new Error("Already connected");
    this._setStatus("creating-session");

    const session = await this._createSession();
    this.dispatchEvent(new CustomEvent("session", { detail: { info: session } }));

    this._setStatus("connecting");

    const pc = new RTCPeerConnection({
      iceServers: [{ urls: "stun:stun.l.google.com:19302" }],
    });
    this._pc = pc;

    pc.addTransceiver(this.options.inputTrack, {
      direction: "sendrecv",
      streams: [new MediaStream([this.options.inputTrack])],
    });

    pc.ontrack = (event) => {
      const [stream] = event.streams;
      const track = event.track.kind === "audio" ? event.track : stream?.getAudioTracks()[0];
      if (!track) return;
      try {
        // Non-standard but widely shipped: ask the browser to buffer a bit
        // of playout to absorb jitter spikes on flaky links.
        /** @type {any} */ (event.receiver).playoutDelayHint = DEFAULT_PLAYOUT_DELAY_S;
      } catch {
        // older browsers ignore the hint
      }
      this.dispatchEvent(new CustomEvent("outputTrack", { detail: { track } }));
    };

    pc.oniceconnectionstatechange = () => {
      const s = pc.iceConnectionState;
      if (s === "connected" || s === "completed") {
        this._clearIceGrace();
        return;
      }
      if (s === "failed") {
        this._clearIceGrace();
        this.dispatchEvent(new CustomEvent("error", { detail: { error: new Error("ICE failed") } }));
        this._setStatus("error");
        return;
      }
      if (s === "disconnected") {
        if (this._iceGraceTimer !== null) return;
        if (typeof document !== "undefined" && document.hidden) {
          this._armIceGraceOnVisibility();
          return;
        }
        this._scheduleIceGrace();
      }
    };

    const dc = pc.createDataChannel("oai-events");
    this._dc = dc;

    dc.addEventListener("open", () => {
      this._setStatus("connected");
      this._sendEvent({
        type: "session.update",
        session: {
          modalities: ["audio", "text"],
          voice: this.options.voice,
          instructions: this.options.instructions,
          input_audio_transcription: { model: "whisper-1" },
          turn_detection: {
            type: "server_vad",
            threshold: 0.5,
            prefix_padding_ms: 300,
            silence_duration_ms: 500,
            create_response: true,
            interrupt_response: true,
          },
        },
      });
    });

    dc.addEventListener("message", (event) => this._handleEvent(event.data));
    dc.addEventListener("error", (event) => {
      console.error("[s2s] data channel error:", event);
    });

    console.log("[s2s] createOffer…");
    const rawOffer = await pc.createOffer();
    const offerSdp = patchOpusFmtp(rawOffer.sdp ?? "");
    console.log("[s2s] setLocalDescription (offer, %d chars)", offerSdp.length);
    await pc.setLocalDescription({ type: "offer", sdp: offerSdp });

    const httpUrl = wssToHttps(session.connectUrl);
    console.log("[s2s] POST SDP offer ->", httpUrl);

    let response;
    try {
      response = await fetch(httpUrl, {
        method: "POST",
        headers: { "Content-Type": "application/sdp" },
        body: offerSdp,
      });
    } catch (err) {
      const detail = err instanceof Error ? `${err.name}: ${err.message}` : String(err);
      throw new Error(`SDP exchange fetch failed (${detail})`);
    }
    console.log("[s2s] SDP exchange status=%d", response.status);

    if (!response.ok) {
      const text = await response.text().catch(() => "");
      throw new Error(`SDP handshake failed (${response.status}): ${text}`);
    }

    const answerSdp = patchOpusFmtp(await response.text());
    await pc.setRemoteDescription({ type: "answer", sdp: answerSdp });
    console.log("[s2s] SDP exchange complete");
  }

  /**
   * @returns {Promise<S2sSessionInfo>}
   */
  async _createSession() {
    const url = `${trimTrailingSlash(this.options.loadBalancerUrl)}/session`;
    console.log("[s2s] POST", url);
    const response = await fetch(url, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: "{}",
    });
    if (!response.ok) {
      const text = await response.text().catch(() => "");
      throw new Error(`/session failed (${response.status}): ${text}`);
    }
    const json = await response.json();
    return {
      sessionId: json.session_id,
      connectUrl: json.connect_url,
      websocketUrl: json.websocket_url,
      sessionToken: json.session_token,
      pendingTimeoutS: json.pending_timeout_s,
    };
  }

  _clearIceGrace() {
    if (this._iceGraceTimer !== null) {
      clearTimeout(this._iceGraceTimer);
      this._iceGraceTimer = null;
    }
    if (this._pendingVisibilityHandler) {
      document.removeEventListener("visibilitychange", this._pendingVisibilityHandler);
      this._pendingVisibilityHandler = null;
    }
  }

  _scheduleIceGrace() {
    if (this._iceGraceTimer !== null) return;
    console.warn("[s2s] ICE disconnected, grace=%dms", ICE_DISCONNECT_GRACE_MS);
    this._iceGraceTimer = window.setTimeout(() => {
      this._iceGraceTimer = null;
      if (!this._pc) return;
      const still = this._pc.iceConnectionState;
      if (still === "disconnected" || still === "failed") {
        this.dispatchEvent(
          new CustomEvent("error", {
            detail: { error: new Error(`ICE stuck in '${still}' for > ${ICE_DISCONNECT_GRACE_MS}ms`) },
          }),
        );
        this._setStatus("error");
      }
    }, ICE_DISCONNECT_GRACE_MS);
  }

  _armIceGraceOnVisibility() {
    if (this._pendingVisibilityHandler) return;
    const handler = () => {
      if (document.hidden) return;
      document.removeEventListener("visibilitychange", handler);
      this._pendingVisibilityHandler = null;
      if (!this._pc) return;
      const s = this._pc.iceConnectionState;
      if (s === "connected" || s === "completed") return;
      if (s === "failed") {
        this.dispatchEvent(new CustomEvent("error", { detail: { error: new Error("ICE failed") } }));
        this._setStatus("error");
        return;
      }
      this._scheduleIceGrace();
    };
    document.addEventListener("visibilitychange", handler);
    this._pendingVisibilityHandler = handler;
  }

  /** @param {string} raw */
  _handleEvent(raw) {
    /** @type {{ type?: string; [k: string]: any }} */
    let event;
    try {
      event = JSON.parse(raw);
    } catch {
      return;
    }

    if (typeof event.type === "string") {
      console.debug("[s2s] event:", event.type);
    }

    switch (event.type) {
      case "input_audio_buffer.speech_started":
        this._setStatus("user-speaking");
        break;

      case "input_audio_buffer.speech_stopped":
        if (this._status === "user-speaking") {
          this._setStatus("processing");
        }
        break;

      case "response.audio.delta":
      case "response.output_audio.delta":
        this._markAudible();
        break;

      case "response.content_part.added": {
        const part = event.part;
        if (part?.type === "audio" || part?.type === "output_audio") {
          this._markAudible();
        }
        break;
      }

      case "response.created":
      case "response.output_item.added":
        if (this._status === "connected" || this._status === "user-speaking") {
          this._setStatus("processing");
        }
        break;

      case "response.done":
      case "response.cancelled":
        if (this._status === "ai-speaking" || this._status === "processing") {
          this._setStatus("connected");
        }
        break;

      case "conversation.item.input_audio_transcription.delta": {
        const delta = typeof event.delta === "string" ? event.delta : "";
        if (delta) {
          this.dispatchEvent(
            new CustomEvent("transcript", {
              detail: /** @type {TranscriptEvent} */ ({ role: "user", text: delta, partial: true }),
            }),
          );
        }
        break;
      }

      case "conversation.item.input_audio_transcription.completed": {
        const transcript = typeof event.transcript === "string" ? event.transcript : "";
        if (transcript) {
          this.dispatchEvent(
            new CustomEvent("transcript", {
              detail: /** @type {TranscriptEvent} */ ({ role: "user", text: transcript, partial: false }),
            }),
          );
        }
        break;
      }

      case "response.audio_transcript.delta":
      case "response.output_audio_transcript.delta": {
        this._markAudible();
        const delta = typeof event.delta === "string" ? event.delta : "";
        if (delta) {
          this.dispatchEvent(
            new CustomEvent("transcript", {
              detail: /** @type {TranscriptEvent} */ ({ role: "assistant", text: delta, partial: true }),
            }),
          );
        }
        break;
      }

      case "response.audio_transcript.done":
      case "response.output_audio_transcript.done": {
        const transcript = typeof event.transcript === "string" ? event.transcript : "";
        if (transcript) {
          this.dispatchEvent(
            new CustomEvent("transcript", {
              detail: /** @type {TranscriptEvent} */ ({ role: "assistant", text: transcript, partial: false }),
            }),
          );
        }
        break;
      }

      case "error": {
        const err = event.error;
        this.dispatchEvent(
          new CustomEvent("error", { detail: { error: new Error(err?.message ?? "Server error") } }),
        );
        break;
      }
    }
  }

  /** @param {Record<string, unknown>} event */
  _sendEvent(event) {
    if (!this._dc || this._dc.readyState !== "open") return;
    this._dc.send(JSON.stringify(event));
  }

  /** Update the live session (voice / instructions / etc.) without reconnecting. */
  /** @param {{ voice?: string; instructions?: string }} patch */
  updateSession(patch) {
    /** @type {Record<string, unknown>} */
    const session = { modalities: ["audio", "text"] };
    if (patch.voice) session.voice = patch.voice;
    if (patch.instructions) session.instructions = patch.instructions;
    this._sendEvent({ type: "session.update", session });
  }

  /** @returns {Promise<void>} */
  async close() {
    this._clearIceGrace();
    try {
      this._dc?.close();
    } catch {
      // ignored
    }
    this._dc = null;
    try {
      this._pc?.close();
    } catch {
      // ignored
    }
    this._pc = null;
    this._setStatus("closed");
  }
}

/** @param {string} url */
function trimTrailingSlash(url) {
  return url.endsWith("/") ? url.slice(0, -1) : url;
}

/** @param {string} url */
function wssToHttps(url) {
  if (url.startsWith("wss://")) return "https://" + url.slice("wss://".length);
  if (url.startsWith("ws://")) return "http://" + url.slice("ws://".length);
  return url;
}

/**
 * Rewrite each `a=fmtp:<pt> ...` line for Opus payload types so it carries
 * our preferred parameters (FEC on, DTX off). Missing fmtp lines are
 * synthesized next to the matching `a=rtpmap`.
 *
 * @param {string} sdp
 * @returns {string}
 */
function patchOpusFmtp(sdp) {
  if (!sdp) return sdp;

  const rtpmapRe = /^a=rtpmap:(\d+)\s+opus\/48000/gim;
  /** @type {string[]} */
  const payloadTypes = [];
  let m;
  while ((m = rtpmapRe.exec(sdp)) !== null) {
    payloadTypes.push(m[1]);
  }
  if (payloadTypes.length === 0) return sdp;

  let patched = sdp;
  for (const pt of payloadTypes) {
    const fmtpLineRe = new RegExp(`^a=fmtp:${pt}\\s+([^\\r\\n]*)`, "m");
    const existing = fmtpLineRe.exec(patched);
    if (existing) {
      const merged = mergeFmtpParams(existing[1], OPUS_FMTP_OVERRIDES);
      patched = patched.replace(fmtpLineRe, `a=fmtp:${pt} ${merged}`);
    } else {
      const merged = mergeFmtpParams("", OPUS_FMTP_OVERRIDES);
      const rtpmapLineRe = new RegExp(`(^a=rtpmap:${pt}\\s+opus/48000[^\\r\\n]*)`, "m");
      patched = patched.replace(rtpmapLineRe, `$1\r\na=fmtp:${pt} ${merged}`);
    }
  }
  return patched;
}

/**
 * @param {string} existing
 * @param {Record<string, string>} overrides
 * @returns {string}
 */
function mergeFmtpParams(existing, overrides) {
  /** @type {Map<string, string>} */
  const params = new Map();
  for (const chunk of existing.split(";")) {
    const kv = chunk.trim();
    if (!kv) continue;
    const eq = kv.indexOf("=");
    if (eq < 0) params.set(kv, "");
    else params.set(kv.slice(0, eq).trim(), kv.slice(eq + 1).trim());
  }
  for (const [k, v] of Object.entries(overrides)) {
    params.set(k, v);
  }
  return Array.from(params.entries())
    .map(([k, v]) => (v ? `${k}=${v}` : k))
    .join(";");
}