refactor(api): make WebSocket proxy generic over STT providers (#2080)

* refactor(api): make WebSocket proxy generic over STT providers

- Rename DeepgramProxyConnection to WsProxyConnection
- Add WsProxyOptions for configurable headers and control message types
- Add x-owh-upstream-url and x-owh-upstream-auth headers for Rust client to specify upstream
- Add createProxyFromRequest factory that reads headers to determine upstream
- Keep createDeepgramProxy and buildDeepgramUrl for backwards compatibility
- Update listen.ts to use the new generic createProxyFromRequest

Co-Authored-By: yujonglee <[email protected]>

* fix(api): wrap createProxyFromRequest in try block for URL parsing errors

- Move createProxyFromRequest call inside try block to catch URL parsing errors
- Change error string from 'upstream_unavailable' to 'upstream_connect_failed'

Co-Authored-By: yujonglee <[email protected]>

* rename

* split files

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Author: yujonglee

apps/api/src/env.ts

@@ -14,6 +14,8 @@ export const env = createEnv({
     STRIPE_WEBHOOK_SECRET: z.string().min(1),
     OPENROUTER_API_KEY: z.string().min(1),
     DEEPGRAM_API_KEY: z.string().min(1),
+    ASSEMBLYAI_API_KEY: z.string().min(1).optional(),
+    SONIOX_API_KEY: z.string().min(1).optional(),
   },
   runtimeEnv: Bun.env,
   emptyStringAsUndefined: true,

apps/api/src/listen.ts

@@ -2,24 +2,24 @@ import type { Handler } from "hono";
 import { upgradeWebSocket } from "hono/bun";
 
 import {
-  buildDeepgramUrl,
-  DeepgramProxyConnection,
+  createProxyFromRequest,
   normalizeWsData,
-} from "./deepgram";
+  WsProxyConnection,
+} from "./stt";
 
 export const listenSocketHandler: Handler = async (c, next) => {
   const clientUrl = new URL(c.req.url, "http://localhost");
-  const deepgramUrl = buildDeepgramUrl(clientUrl).toString();
 
-  const connection = new DeepgramProxyConnection(deepgramUrl);
+  let connection: WsProxyConnection;
   try {
+    connection = createProxyFromRequest(clientUrl, c.req.raw.headers);
     await connection.preconnectUpstream();
   } catch (error) {
     console.error("Failed to establish upstream connection", error);
     const detail =
-      error instanceof Error ? error.message : "upstream_unavailable";
+      error instanceof Error ? error.message : "upstream_connect_failed";
     const status = detail === "upstream_connect_timeout" ? 504 : 502;
-    return c.json({ error: "upstream_unavailable", detail }, status);
+    return c.json({ error: "upstream_connect_failed", detail }, status);
   }
 
   const handler = upgradeWebSocket(() => {

apps/api/src/stt/assemblyai.ts

@@ -0,0 +1,31 @@
+import { env } from "../env";
+import { WsProxyConnection } from "./connection";
+
+const CONTROL_MESSAGE_TYPES = new Set(["Terminate"]);
+
+export const buildAssemblyAIUrl = (incomingUrl: URL) => {
+  const target = new URL("wss://streaming.assemblyai.com/v3/ws");
+
+  incomingUrl.searchParams.forEach((value, key) => {
+    if (key !== "provider") {
+      target.searchParams.set(key, value);
+    }
+  });
+
+  return target;
+};
+
+export const createAssemblyAIProxy = (incomingUrl: URL): WsProxyConnection => {
+  const apiKey = env.ASSEMBLYAI_API_KEY;
+  if (!apiKey) {
+    throw new Error("ASSEMBLYAI_API_KEY not configured");
+  }
+
+  const target = buildAssemblyAIUrl(incomingUrl);
+  return new WsProxyConnection(target.toString(), {
+    headers: {
+      authorization: apiKey,
+    },
+    controlMessageTypes: CONTROL_MESSAGE_TYPES,
+  });
+};

apps/api/src/deepgram.ts apps/api/src/stt/connection.tsapps/api/src/deepgram.ts renamed to apps/api/src/stt/connection.ts

@@ -1,17 +1,27 @@
 import type { ServerWebSocket, WebSocketOptions } from "bun";
 
-import { env } from "./env";
+import {
+  getPayloadSize,
+  normalizeWsData,
+  payloadIsControlMessage,
+  type WsPayload,
+} from "./utils";
 
 const DEFAULT_CLOSE_CODE = 1011;
 const UPSTREAM_ERROR_TIMEOUT = 1000;
 const UPSTREAM_CONNECT_TIMEOUT = 5000;
 const MAX_PENDING_QUEUE_BYTES = 5 * 1024 * 1024; // 5 MiB
-const TEXT_ENCODER = new TextEncoder();
 
 type QueuedPayload = { payload: WsPayload; size: number };
 
-// https://developers.deepgram.com/docs/lower-level-websockets
-export class DeepgramProxyConnection {
+export type WsProxyOptions = {
+  headers?: Record<string, string>;
+  controlMessageTypes?: ReadonlySet<string>;
+};
+
+const DEFAULT_CONTROL_MESSAGE_TYPES = new Set<string>();
+
+export class WsProxyConnection {
   private upstream?: InstanceType<typeof WebSocket>;
   private upstreamReady = false;
   private shuttingDown = false;
@@ -25,7 +35,17 @@ export class DeepgramProxyConnection {
   private upstreamReadyResolve: (() => void) | null = null;
   private upstreamReadyReject: ((error: Error) => void) | null = null;
 
-  constructor(private deepgramUrl: string) {}
+  private readonly headers?: Record<string, string>;
+  private readonly controlMessageTypes: ReadonlySet<string>;
+
+  constructor(
+    private upstreamUrl: string,
+    options: WsProxyOptions = {},
+  ) {
+    this.headers = options.headers;
+    this.controlMessageTypes =
+      options.controlMessageTypes ?? DEFAULT_CONTROL_MESSAGE_TYPES;
+  }
 
   private clearErrorTimeout() {
     if (this.upstreamErrorTimeout) {
@@ -80,18 +100,17 @@ export class DeepgramProxyConnection {
       return;
     }
 
-    const wsOptions: WebSocketOptions = {
-      headers: {
-        Authorization: `Token ${env.DEEPGRAM_API_KEY}`,
-      },
-    };
+    const wsOptions: WebSocketOptions =
+      this.headers && Object.keys(this.headers).length > 0
+        ? { headers: this.headers }
+        : {};
 
     this.upstream = new (globalThis.WebSocket as {
       new (
         url: string | URL,
         options?: WebSocketOptions,
       ): InstanceType<typeof WebSocket>;
-    })(this.deepgramUrl, wsOptions);
+    })(this.upstreamUrl, wsOptions);
 
     this.upstream.binaryType = "arraybuffer";
     this.setupUpstreamHandlers();
@@ -324,7 +343,10 @@ export class DeepgramProxyConnection {
       return;
     }
 
-    const isControlPayload = payloadIsControlMessage(payload);
+    const isControlPayload = payloadIsControlMessage(
+      payload,
+      this.controlMessageTypes,
+    );
 
     if (!this.upstreamReady) {
       this.enqueuePendingPayload(payload, isControlPayload);
@@ -360,85 +382,3 @@ export class DeepgramProxyConnection {
     this.pendingBytes += size;
   }
 }
-
-export const buildDeepgramUrl = (incomingUrl: URL) => {
-  const target = new URL("wss://api.deepgram.com/v1/listen");
-
-  incomingUrl.searchParams.forEach((value, key) => {
-    target.searchParams.set(key, value);
-  });
-  target.searchParams.set("model", "nova-3-general");
-  target.searchParams.set("mip_opt_out", "false");
-
-  return target;
-};
-
-export type WsPayload = string | Uint8Array;
-
-// https://bun.com/docs/runtime/http/websockets
-export const normalizeWsData = async (
-  data: unknown,
-): Promise<WsPayload | null> => {
-  if (typeof data === "string") {
-    return data;
-  }
-
-  if (data instanceof Uint8Array) {
-    return cloneBinaryPayload(data);
-  }
-
-  if (data instanceof ArrayBuffer) {
-    return new Uint8Array(data);
-  }
-
-  if (ArrayBuffer.isView(data)) {
-    return cloneBinaryPayload(data);
-  }
-
-  if (typeof Blob !== "undefined" && data instanceof Blob) {
-    const arrayBuffer = await data.arrayBuffer();
-    return new Uint8Array(arrayBuffer);
-  }
-
-  return null;
-};
-
-const cloneBinaryPayload = (input: ArrayBuffer | ArrayBufferView) => {
-  const view =
-    input instanceof ArrayBuffer
-      ? new Uint8Array(input)
-      : new Uint8Array(input.buffer, input.byteOffset, input.byteLength);
-  const copy = new Uint8Array(view.byteLength);
-  copy.set(view);
-  return copy;
-};
-
-const getPayloadSize = (payload: WsPayload) => {
-  if (typeof payload === "string") {
-    return TEXT_ENCODER.encode(payload).byteLength;
-  }
-  return payload.byteLength;
-};
-
-const CONTROL_MESSAGE_TYPES = new Set(["KeepAlive", "CloseStream", "Finalize"]);
-
-const payloadIsControlMessage = (payload: WsPayload) => {
-  if (typeof payload !== "string") {
-    return false;
-  }
-
-  try {
-    const parsed = JSON.parse(payload);
-    if (
-      typeof parsed === "object" &&
-      parsed !== null &&
-      CONTROL_MESSAGE_TYPES.has(parsed.type)
-    ) {
-      return true;
-    }
-  } catch {
-    // ignore parse errors
-  }
-
-  return false;
-};

apps/api/src/stt/deepgram.ts

@@ -0,0 +1,28 @@
+import { env } from "../env";
+import { WsProxyConnection } from "./connection";
+
+const CONTROL_MESSAGE_TYPES = new Set(["KeepAlive", "CloseStream", "Finalize"]);
+
+export const buildDeepgramUrl = (incomingUrl: URL) => {
+  const target = new URL("wss://api.deepgram.com/v1/listen");
+
+  incomingUrl.searchParams.forEach((value, key) => {
+    if (key !== "provider") {
+      target.searchParams.set(key, value);
+    }
+  });
+  target.searchParams.set("model", "nova-3-general");
+  target.searchParams.set("mip_opt_out", "false");
+
+  return target;
+};
+
+export const createDeepgramProxy = (incomingUrl: URL): WsProxyConnection => {
+  const target = buildDeepgramUrl(incomingUrl);
+  return new WsProxyConnection(target.toString(), {
+    headers: {
+      Authorization: `Token ${env.DEEPGRAM_API_KEY}`,
+    },
+    controlMessageTypes: CONTROL_MESSAGE_TYPES,
+  });
+};

apps/api/src/stt/index.ts

@@ -0,0 +1,46 @@
+import { createAssemblyAIProxy } from "./assemblyai";
+import { WsProxyConnection } from "./connection";
+import { createDeepgramProxy } from "./deepgram";
+import { createSonioxProxy } from "./soniox";
+
+export { WsProxyConnection, type WsProxyOptions } from "./connection";
+export { normalizeWsData, type WsPayload } from "./utils";
+export { buildDeepgramUrl, createDeepgramProxy } from "./deepgram";
+export { buildAssemblyAIUrl, createAssemblyAIProxy } from "./assemblyai";
+export { buildSonioxUrl, createSonioxProxy } from "./soniox";
+
+export const UPSTREAM_URL_HEADER = "x-owh-upstream-url";
+export const UPSTREAM_AUTH_HEADER = "x-owh-upstream-auth";
+
+export type SttProvider = "deepgram" | "assemblyai" | "soniox";
+
+export function createProxyFromRequest(
+  incomingUrl: URL,
+  reqHeaders: Headers,
+): WsProxyConnection {
+  const upstreamOverride = reqHeaders.get(UPSTREAM_URL_HEADER);
+  const rawAuth = reqHeaders.get(UPSTREAM_AUTH_HEADER);
+
+  if (upstreamOverride) {
+    const url = new URL(upstreamOverride);
+    const headers =
+      rawAuth && rawAuth.length > 0 ? { Authorization: rawAuth } : undefined;
+
+    return new WsProxyConnection(url.toString(), {
+      headers,
+    });
+  }
+
+  const provider =
+    (incomingUrl.searchParams.get("provider") as SttProvider) || "deepgram";
+
+  switch (provider) {
+    case "assemblyai":
+      return createAssemblyAIProxy(incomingUrl);
+    case "soniox":
+      return createSonioxProxy();
+    case "deepgram":
+    default:
+      return createDeepgramProxy(incomingUrl);
+  }
+}

apps/api/src/stt/soniox.ts

@@ -0,0 +1,19 @@
+import { env } from "../env";
+import { WsProxyConnection } from "./connection";
+
+const CONTROL_MESSAGE_TYPES = new Set(["keepalive", "finalize"]);
+
+export const buildSonioxUrl = () => {
+  return new URL("wss://stt-rt.soniox.com/transcribe-websocket");
+};
+
+export const createSonioxProxy = (): WsProxyConnection => {
+  if (!env.SONIOX_API_KEY) {
+    throw new Error("SONIOX_API_KEY not configured");
+  }
+
+  const target = buildSonioxUrl();
+  return new WsProxyConnection(target.toString(), {
+    controlMessageTypes: CONTROL_MESSAGE_TYPES,
+  });
+};