From 61fc15a3dafdfde61a4b61cbb8275e2897c8342f Mon Sep 17 00:00:00 2001 From: Yash Hayaran Date: Wed, 8 Jul 2026 13:02:33 +0530 Subject: [PATCH] Add mint API key support for session authentication in realtime client - Introduced `--mint-api-key` argument for specifying the API key. - Implemented `_mint_auth_headers` function to include the API key in session requests. - Updated session initialization and WebSocket connection methods to utilize the client secret token. - Enhanced error handling for session authentication failures. --- riva/client/argparse_utils.py | 5 ++ riva/client/realtime.py | 111 ++++++++++++++++++++++++----- scripts/tts/realtime_tts_client.py | 6 +- 3 files changed, 102 insertions(+), 20 deletions(-) diff --git a/riva/client/argparse_utils.py b/riva/client/argparse_utils.py index a8cc1a7..0579108 100644 --- a/riva/client/argparse_utils.py +++ b/riva/client/argparse_utils.py @@ -132,4 +132,9 @@ def add_connection_argparse_parameters(parser: argparse.ArgumentParser) -> argpa def add_realtime_config_argparse_parameters(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument("--endpoint", default="/v1/realtime", help="Endpoint to WebSocket server endpoint.") parser.add_argument("--query-params", default="intent=transcription", help="Query parameters to WebSocket server endpoint.") + parser.add_argument( + "--mint-api-key", + default=None, + help="API key for POST /v1/realtime/*_sessions (Authorization: Bearer).", + ) return parser \ No newline at end of file diff --git a/riva/client/realtime.py b/riva/client/realtime.py index 772cd5c..45fcc99 100644 --- a/riva/client/realtime.py +++ b/riva/client/realtime.py @@ -7,7 +7,7 @@ import logging import queue import uuid -from typing import Dict, Any, Generator +from typing import Dict, Any, Generator, Optional import requests import websockets @@ -21,6 +21,69 @@ logger = logging.getLogger(__name__) +def _mint_auth_headers(args: argparse.Namespace) -> Dict[str, str]: + """Build HTTP headers for session mint endpoints, including optional tier-1 auth.""" + headers = {"Content-Type": "application/json"} + api_key = (getattr(args, "mint_api_key", None) or "").strip() + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + return headers + + +def _server_error_detail(response) -> str: + """Extract the FastAPI ``detail`` message from an error response, falling back to raw text.""" + try: + payload = response.json() + except ValueError: + return response.text.strip() + if isinstance(payload, dict) and payload.get("detail"): + return str(payload["detail"]) + return response.text.strip() + + +def _raise_for_session_status(response, provided_key: bool) -> None: + """Raise a clear, actionable error when the session mint request fails.""" + if response.status_code == 200: + return + if response.status_code in (401, 403): + if provided_key: + hint = "The --mint-api-key value was rejected by the server. Verify it matches the server's REALTIME_AUTH_MINT_API_KEY." + else: + hint = "This server requires an API key on the session endpoint. Pass it with --mint-api-key ." + raise Exception( + f"Session authentication failed (HTTP {response.status_code}: {_server_error_detail(response)}). {hint}" + ) + raise Exception( + f"Failed to initialize session. Status: {response.status_code}, Error: {_server_error_detail(response)}" + ) + + +def _extract_client_secret_token(session_data: Dict[str, Any]) -> str: + """Return the ephemeral client_secret token from a session-creation response.""" + client_secret = session_data.get("client_secret") + if not isinstance(client_secret, dict): + raise ValueError("Session response missing client_secret") + token = client_secret.get("value") + if not token: + raise ValueError("Session response missing client_secret.value") + return token + + +def _build_websocket_url( + server: str, + endpoint: str, + query_params: str, + token: str, + use_ssl: bool, +) -> str: + """Build the WebSocket URL with intent and ephemeral token query parameters.""" + query = query_params.strip() + if "token=" not in query: + query = f"{query}&token={token}" if query else f"token={token}" + scheme = "wss" if use_ssl else "ws" + return f"{scheme}://{server}{endpoint}?{query}" + + class RealtimeClientASR: """Client for real-time transcription via WebSocket connection.""" @@ -33,6 +96,7 @@ def __init__(self, args: argparse.Namespace): self.args = args self.websocket = None self.session_config = None + self._client_secret_token: Optional[str] = None # Input audio playback self.input_audio_queue = queue.Queue() @@ -49,6 +113,7 @@ async def connect(self): # Initialize session via HTTP POST session_data = await self._initialize_http_session() self.session_config = session_data + self._client_secret_token = _extract_client_secret_token(session_data) # Connect to WebSocket await self._connect_websocket() @@ -66,7 +131,7 @@ async def connect(self): async def _initialize_http_session(self) -> Dict[str, Any]: """Initialize session via HTTP POST request.""" - headers = {"Content-Type": "application/json"} + headers = _mint_auth_headers(self.args) uri = f"http://{self.args.server}/v1/realtime/transcription_sessions" if self.args.use_ssl: uri = f"https://{self.args.server}/v1/realtime/transcription_sessions" @@ -79,11 +144,7 @@ async def _initialize_http_session(self) -> Dict[str, Any]: verify=self.args.ssl_root_cert if self.args.ssl_root_cert else True ) - if response.status_code != 200: - raise Exception( - f"Failed to initialize session. Status: {response.status_code}, " - f"Error: {response.text}" - ) + _raise_for_session_status(response, provided_key=bool(getattr(self.args, "mint_api_key", None))) session_data = response.json() logger.debug("Session initialized: %s", session_data) @@ -91,11 +152,18 @@ async def _initialize_http_session(self) -> Dict[str, Any]: async def _connect_websocket(self): """Connect to WebSocket endpoint.""" + if not self._client_secret_token: + raise ValueError("client_secret token is required before opening the WebSocket") + ssl_context = None - ws_url = f"ws://{self.args.server}{self.args.endpoint}?{self.args.query_params}" + ws_url = _build_websocket_url( + self.args.server, + self.args.endpoint, + self.args.query_params, + self._client_secret_token, + self.args.use_ssl, + ) if self.args.use_ssl: - ws_url = f"wss://{self.args.server}{self.args.endpoint}?{self.args.query_params}" - ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) # Load a custom CA certificate bundle if self.args.ssl_root_cert: @@ -500,6 +568,7 @@ def __init__(self, args: argparse.Namespace): self.args = args self.websocket = None self.session_config = None + self._client_secret_token: Optional[str] = None self.audio_data = [] self.is_synthesis_complete = False self.wav_file = None # WAV file handle for streaming write @@ -528,6 +597,7 @@ async def connect(self): logger.info("Initializing HTTP session...") session_data = await self._initialize_http_session() self.session_config = session_data + self._client_secret_token = _extract_client_secret_token(session_data) logger.info("HTTP session initialized successfully") # Connect to WebSocket @@ -557,7 +627,7 @@ async def connect(self): async def _initialize_http_session(self) -> Dict[str, Any]: """Initialize session via HTTP POST request.""" - headers = {"Content-Type": "application/json"} + headers = _mint_auth_headers(self.args) uri = f"http://{self.args.server}/v1/realtime/synthesis_sessions" if self.args.use_ssl: uri = f"https://{self.args.server}/v1/realtime/synthesis_sessions" @@ -595,11 +665,7 @@ async def _initialize_http_session(self) -> Dict[str, Any]: logger.error("HTTP request failed: %s", e) raise - if response.status_code != 200: - raise Exception( - f"Failed to initialize session. Status: {response.status_code}, " - f"Error: {response.text}" - ) + _raise_for_session_status(response, provided_key=bool(getattr(self.args, "mint_api_key", None))) session_data = response.json() logger.info("Session initialized: %s", session_data) @@ -607,11 +673,18 @@ async def _initialize_http_session(self) -> Dict[str, Any]: async def _connect_websocket(self): """Connect to WebSocket endpoint.""" + if not self._client_secret_token: + raise ValueError("client_secret token is required before opening the WebSocket") + ssl_context = None - ws_url = f"ws://{self.args.server}{self.args.endpoint}?{self.args.query_params}" + ws_url = _build_websocket_url( + self.args.server, + self.args.endpoint, + self.args.query_params, + self._client_secret_token, + self.args.use_ssl, + ) if self.args.use_ssl: - ws_url = f"wss://{self.args.server}{self.args.endpoint}?{self.args.query_params}" - ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) if self.args.ssl_root_cert: ssl_context.load_verify_locations(self.args.ssl_root_cert) diff --git a/scripts/tts/realtime_tts_client.py b/scripts/tts/realtime_tts_client.py index c3ac728..a4fc764 100644 --- a/scripts/tts/realtime_tts_client.py +++ b/scripts/tts/realtime_tts_client.py @@ -20,7 +20,10 @@ import websockets from websockets.exceptions import WebSocketException -from riva.client.argparse_utils import add_connection_argparse_parameters +from riva.client.argparse_utils import ( + add_connection_argparse_parameters, + add_realtime_config_argparse_parameters, +) from riva.client.audio_io import SoundCallBack from riva.client.realtime import RealtimeClientTTS @@ -144,6 +147,7 @@ def parse_args() -> argparse.Namespace: # Add connection parameters parser = add_connection_argparse_parameters(parser) + parser = add_realtime_config_argparse_parameters(parser) # Override default server for realtime TTS (WebSocket endpoint, not gRPC) parser.set_defaults(server="localhost:9000")