fix(gradium): identify the STT plugin to the API with x-api-source: livekit-plugin

livekit-plugins/livekit-plugins-gradium/livekit/plugins/gradium/api.py

@@ -0,0 +1,8 @@
+# Value of the `x-api-source` header on every Gradium websocket, so sessions opened
+# by this plugin can be told apart from direct API traffic on the same API key.
+API_SOURCE = "livekit-plugin"
+
+
+def ws_headers(api_key: str) -> dict[str, str]:
+    """Headers for a Gradium websocket: authentication plus source attribution."""
+    return {"x-api-key": api_key, "x-api-source": API_SOURCE}

livekit-plugins/livekit-plugins-gradium/livekit/plugins/gradium/stt.py

@@ -38,6 +38,7 @@
 from livekit.agents.utils import AudioBuffer, is_given
 from livekit.rtc import AudioFrame
 
+from .api import ws_headers
 from .log import logger
 
 STTEncoding = Literal["pcm_s16le"]
@@ -411,9 +412,7 @@ async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
                     await ws.close()
 
     async def _connect_ws(self) -> aiohttp.ClientWebSocketResponse:
-        headers = {"x-api-key": self._api_key, "x-api-source": "livekit"}
-
-        ws = await self._session.ws_connect(self._model_endpoint, headers=headers)
+        ws = await self._session.ws_connect(self._model_endpoint, headers=ws_headers(self._api_key))
 
         # Build and send the setup payload as the first message
         setup_msg: dict[str, Any] = {

tests/test_plugin_gradium_api_source.py

@@ -0,0 +1,51 @@
+"""The Gradium STT websocket identifies the plugin to the API through `x-api-source`."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from livekit.agents import DEFAULT_API_CONNECT_OPTIONS
+
+pytestmark = pytest.mark.unit
+
+# Pinned on purpose rather than imported: the value is what the API records, so a
+# rename in the plugin must fail here.
+API_SOURCE = "livekit-plugin"
+
+
+async def test_stt_sends_api_source_header():
+    from livekit.plugins.gradium import STT
+    from livekit.plugins.gradium.stt import SpeechStream, STTOptions
+
+    recorded: dict[str, Any] = {}
+
+    class FakeWebSocket:
+        async def send_str(self, message: str) -> None:
+            pass
+
+    class FakeSession:
+        async def ws_connect(self, url: str, **kwargs: Any) -> FakeWebSocket:
+            recorded["headers"] = kwargs["headers"]
+            return FakeWebSocket()
+
+    def _fake_create_task(coro, *args, **kwargs):
+        coro.close()
+        return MagicMock()
+
+    stt = STT(api_key="test-key")
+    with patch("livekit.agents.stt.stt.asyncio.create_task", side_effect=_fake_create_task):
+        stream = SpeechStream(
+            stt=stt,
+            opts=STTOptions(),
+            conn_options=DEFAULT_API_CONNECT_OPTIONS,
+            api_key="test-key",
+            model_endpoint="wss://api.gradium.ai/api/speech/asr",
+            model_name="default",
+            http_session=FakeSession(),  # type: ignore[arg-type]
+        )
+    await stream._connect_ws()
+
+    assert recorded["headers"] == {"x-api-key": "test-key", "x-api-source": API_SOURCE}