diff --git a/src/agents/models/_openai_websocket.py b/src/agents/models/_openai_websocket.py new file mode 100644 index 0000000000..1687bfc253 --- /dev/null +++ b/src/agents/models/_openai_websocket.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import httpx2 +from openai import AsyncOpenAI, NotGiven, Omit + +from .._httpx_compat import is_legacy_httpx_instance +from ..exceptions import UserError + + +def _is_openai_omitted_value(value: Any) -> bool: + return isinstance(value, Omit | NotGiven) + + +async def refresh_openai_client_api_key_if_supported(client: Any) -> None: + """Refresh dynamic OpenAI client credentials before materializing handshake headers.""" + refresh_api_key = getattr(client, "_refresh_api_key", None) + if callable(refresh_api_key): + await refresh_api_key() + + +def _remove_header(headers: dict[str, str], key: object) -> None: + header_key = str(key) + for existing_key in list(headers): + if existing_key.lower() == header_key.lower(): + del headers[existing_key] + + +def _set_header(headers: dict[str, str], key: object, value: object) -> None: + header_key = str(key) + _remove_header(headers, header_key) + headers[header_key] = str(value) + + +def merge_openai_client_websocket_headers( + client: AsyncOpenAI, + *, + extra_headers: Mapping[str, Any] | None = None, +) -> dict[str, str]: + """Materialize OpenAI client auth/default headers for a WebSocket handshake.""" + headers: dict[str, str] = {} + for source in ( + getattr(client, "auth_headers", {}), + getattr(client, "default_headers", {}), + ): + for key, value in source.items(): + if isinstance(value, NotGiven): + continue + if isinstance(value, Omit): + _remove_header(headers, key) + continue + _set_header(headers, key, value) + + for key, value in (extra_headers or {}).items(): + if isinstance(value, NotGiven): + continue + _remove_header(headers, key) + if isinstance(value, Omit): + continue + headers[str(key)] = str(value) + + return headers + + +def _merge_query_values(params: dict[str, Any], values: Mapping[str, Any]) -> None: + for key, value in values.items(): + query_key = str(key) + if isinstance(value, Omit): + params.pop(query_key, None) + continue + if isinstance(value, NotGiven): + continue + params[query_key] = value + + +def prepare_openai_client_websocket_base_url( + client: AsyncOpenAI, + *, + extra_query: Any = None, + context: str, +) -> httpx2.URL: + """Build the client-derived WebSocket base URL and normalized query parameters. + + Endpoint suffixes and transport-specific fixed query parameters are intentionally left to + each caller. + """ + websocket_base_url = getattr(client, "websocket_base_url", None) + if websocket_base_url is not None: + if is_legacy_httpx_instance(websocket_base_url, "URL"): + websocket_base_url = str(websocket_base_url) + base_url = httpx2.URL(websocket_base_url) + else: + client_base_url = client.base_url + if is_legacy_httpx_instance(client_base_url, "URL"): + client_base_url = str(client_base_url) + base_url = httpx2.URL(client_base_url) + + ws_scheme = {"http": "ws", "https": "wss"}.get(base_url.scheme, base_url.scheme) + base_url = base_url.copy_with(scheme=ws_scheme) + params: dict[str, Any] = dict(base_url.params) + + default_query = getattr(client, "default_query", None) + if default_query is not None and not _is_openai_omitted_value(default_query): + if not isinstance(default_query, Mapping): + raise UserError(f"{context} client default_query must be a mapping.") + _merge_query_values(params, default_query) + + if extra_query is not None and not _is_openai_omitted_value(extra_query): + if not isinstance(extra_query, Mapping): + raise UserError(f"{context} extra_query must be a mapping.") + _merge_query_values(params, extra_query) + + return base_url.copy_with(params=params) diff --git a/src/agents/models/openai_responses.py b/src/agents/models/openai_responses.py index f6ff8e2b3c..d7244e324e 100644 --- a/src/agents/models/openai_responses.py +++ b/src/agents/models/openai_responses.py @@ -91,6 +91,11 @@ from ..util._json import _to_dump_compatible from ..version import __version__ from ._openai_retry import get_openai_retry_advice +from ._openai_websocket import ( + merge_openai_client_websocket_headers, + prepare_openai_client_websocket_base_url, + refresh_openai_client_api_key_if_supported, +) from ._response_terminal import response_error_event_failure_error, response_terminal_failure_error from ._retry_runtime import ( should_disable_provider_managed_retries, @@ -177,10 +182,8 @@ def _materialize_responses_tool_params( async def _refresh_openai_client_api_key_if_supported(client: Any) -> None: - """Refresh client auth if the current OpenAI SDK exposes a refresh hook.""" - refresh_api_key = getattr(client, "_refresh_api_key", None) - if callable(refresh_api_key): - await refresh_api_key() + """Backward-compatible wrapper around shared WebSocket client credential refresh.""" + await refresh_openai_client_api_key_if_supported(client) def _construct_response_stream_event_from_payload( @@ -1533,76 +1536,19 @@ async def _prepare_websocket_request( return frame, ws_url, handshake_headers def _merge_websocket_headers(self, extra_headers: Mapping[str, Any]) -> dict[str, str]: - headers: dict[str, str] = {} - for source in ( - getattr(self._client, "auth_headers", {}), - self._client.default_headers, - ): - for key, value in source.items(): - if _is_openai_omitted_value(value): - continue - header_key = str(key) - for existing_key in list(headers): - if existing_key.lower() == header_key.lower(): - del headers[existing_key] - headers[header_key] = str(value) - - for key, value in extra_headers.items(): - if isinstance(value, NotGiven): - continue - header_key = str(key) - for existing_key in list(headers): - if existing_key.lower() == header_key.lower(): - del headers[existing_key] - if isinstance(value, Omit): - continue - headers[header_key] = str(value) - - return headers + return merge_openai_client_websocket_headers( + self._client, + extra_headers=extra_headers, + ) def _prepare_websocket_url(self, extra_query: Any) -> str: - if self._client.websocket_base_url is not None: - websocket_base_url = self._client.websocket_base_url - if is_legacy_httpx_instance(websocket_base_url, "URL"): - websocket_base_url = str(websocket_base_url) - base_url = httpx2.URL(websocket_base_url) - ws_scheme = {"http": "ws", "https": "wss"}.get(base_url.scheme, base_url.scheme) - base_url = base_url.copy_with(scheme=ws_scheme) - else: - client_base_url = self._client.base_url - ws_scheme = {"http": "ws", "https": "wss"}.get( - client_base_url.scheme, client_base_url.scheme - ) - base_url = client_base_url.copy_with(scheme=ws_scheme) - - params: dict[str, Any] = dict(base_url.params) - default_query = getattr(self._client, "default_query", None) - if default_query is not None and not _is_openai_omitted_value(default_query): - if not isinstance(default_query, Mapping): - raise UserError("Responses websocket client default_query must be a mapping.") - for key, value in default_query.items(): - query_key = str(key) - if isinstance(value, Omit): - params.pop(query_key, None) - continue - if isinstance(value, NotGiven): - continue - params[query_key] = value - - if extra_query is not None and not _is_openai_omitted_value(extra_query): - if not isinstance(extra_query, Mapping): - raise UserError("Responses websocket extra_query must be a mapping.") - for key, value in extra_query.items(): - query_key = str(key) - if isinstance(value, Omit): - params.pop(query_key, None) - continue - if isinstance(value, NotGiven): - continue - params[query_key] = value - + base_url = prepare_openai_client_websocket_base_url( + self._client, + extra_query=extra_query, + context="Responses websocket", + ) path = base_url.path.rstrip("/") + "/responses" - return str(base_url.copy_with(path=path, params=params)) + return str(base_url.copy_with(path=path)) async def _ensure_websocket_connection( self, diff --git a/src/agents/voice/models/openai_stt.py b/src/agents/voice/models/openai_stt.py index cf504d8892..ff47496399 100644 --- a/src/agents/voice/models/openai_stt.py +++ b/src/agents/voice/models/openai_stt.py @@ -13,6 +13,11 @@ from ... import _debug from ...exceptions import AgentsException, UserError from ...logger import logger +from ...models._openai_websocket import ( + merge_openai_client_websocket_headers, + prepare_openai_client_websocket_base_url, + refresh_openai_client_api_key_if_supported, +) from ...tracing import Span, SpanError, TranscriptionSpanData, transcription_span from ...util._error_tracing import get_trace_error from ..exceptions import STTWebsocketConnectionError @@ -58,6 +63,24 @@ def _audio_buffer_to_base64(buffer: npt.NDArray[np.int16 | np.float32]) -> str: return base64.b64encode(buffer.tobytes()).decode("utf-8") +def _prepare_websocket_url(client: AsyncOpenAI) -> str: + base_url = prepare_openai_client_websocket_base_url( + client, + context="Streamed STT websocket", + ) + params: dict[str, Any] = dict(base_url.params) + params["intent"] = "transcription" + path = base_url.path.rstrip("/") + "/realtime" + return str(base_url.copy_with(path=path, params=params)) + + +def _prepare_websocket_headers(client: AsyncOpenAI) -> dict[str, str]: + return merge_openai_client_websocket_headers( + client, + extra_headers={"OpenAI-Log-Session": "1"}, + ) + + async def _wait_for_event( event_queue: asyncio.Queue[dict[str, Any] | ErrorSentinel], expected_types: list[str], @@ -303,12 +326,10 @@ async def _stream_audio( async def _process_websocket_connection(self) -> None: try: + await refresh_openai_client_api_key_if_supported(self._client) async with websockets.connect( - "wss://api.openai.com/v1/realtime?intent=transcription", - additional_headers={ - "Authorization": f"Bearer {self._client.api_key}", - "OpenAI-Log-Session": "1", - }, + _prepare_websocket_url(self._client), + additional_headers=_prepare_websocket_headers(self._client), ) as ws: await self._setup_connection(ws) self._process_events_task = asyncio.create_task(self._handle_events()) diff --git a/tests/voice/test_openai_stt.py b/tests/voice/test_openai_stt.py index 50daf0c2ba..a64c7079f5 100644 --- a/tests/voice/test_openai_stt.py +++ b/tests/voice/test_openai_stt.py @@ -9,9 +9,11 @@ from typing import cast from unittest.mock import AsyncMock, MagicMock, patch +import httpx2 import numpy as np import numpy.typing as npt import pytest +from openai import AsyncOpenAI import agents._debug as _debug from agents import trace @@ -55,6 +57,17 @@ def create_mock_websocket(messages: list[str]) -> AsyncMock: return mock_ws +def create_mock_openai_client(api_key: str = "FAKE_KEY") -> AsyncOpenAI: + client = AsyncMock(api_key=api_key) + client.websocket_base_url = None + client.base_url = httpx2.URL("https://api.openai.com/v1/") + client.default_query = {} + client.auth_headers = {"Authorization": f"Bearer {api_key}"} + client.default_headers = {} + client._refresh_api_key = AsyncMock() + return cast(AsyncOpenAI, client) + + def fake_time(increment: int): current = 1000 while True: @@ -67,7 +80,7 @@ def fake_time(increment: int): async def test_transcribe_turns_propagates_consumer_cancellation(monkeypatch) -> None: session = OpenAISTTTranscriptionSession( input=StreamedAudioInput(), - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=STTModelSettings(), trace_include_sensitive_data=False, @@ -105,7 +118,7 @@ async def hold_connection_open() -> None: async def test_transcribe_turns_closes_owned_tasks_after_yield(monkeypatch) -> None: session = OpenAISTTTranscriptionSession( input=StreamedAudioInput(), - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=STTModelSettings(), trace_include_sensitive_data=False, @@ -165,7 +178,7 @@ async def hold_connection_open() -> None: async def test_close_finishes_span_started_while_websocket_close_is_pending() -> None: session = OpenAISTTTranscriptionSession( input=StreamedAudioInput(), - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=STTModelSettings(), trace_include_sensitive_data=False, @@ -223,7 +236,7 @@ async def test_transcribe_turns_preserves_consumer_exception_when_cleanup_fails( ) -> None: session = OpenAISTTTranscriptionSession( input=StreamedAudioInput(), - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=STTModelSettings(), trace_include_sensitive_data=False, @@ -270,7 +283,7 @@ async def fail_cleanup() -> None: async def test_transcribe_turns_propagates_cancellation_during_cleanup(monkeypatch) -> None: session = OpenAISTTTranscriptionSession( input=StreamedAudioInput(), - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=STTModelSettings(), trace_include_sensitive_data=False, @@ -307,7 +320,7 @@ async def test_transcribe_turns_preserves_terminal_error_when_close_fails( ) -> None: session = OpenAISTTTranscriptionSession( input=StreamedAudioInput(), - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=STTModelSettings(), trace_include_sensitive_data=False, @@ -372,7 +385,7 @@ async def test_non_json_messages_should_crash(): session = OpenAISTTTranscriptionSession( input=input_audio, - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=stt_settings, trace_include_sensitive_data=False, @@ -412,7 +425,7 @@ async def test_session_connects_and_configures_successfully(): session = OpenAISTTTranscriptionSession( input=input_audio, - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=stt_settings, trace_include_sensitive_data=False, @@ -472,7 +485,7 @@ async def test_stream_audio_sends_pcm16( session = OpenAISTTTranscriptionSession( input=audio_input, - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=stt_settings, trace_include_sensitive_data=False, @@ -548,7 +561,7 @@ async def test_transcription_event_puts_output_in_queue(created, updated, comple session = OpenAISTTTranscriptionSession( input=audio_input, - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=stt_settings, trace_include_sensitive_data=False, @@ -594,7 +607,7 @@ def fake_time_func(): session = OpenAISTTTranscriptionSession( input=audio_input, - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=stt_settings, trace_include_sensitive_data=False, @@ -643,7 +656,7 @@ async def test_session_error_event(monkeypatch: pytest.MonkeyPatch): session = OpenAISTTTranscriptionSession( input=audio_input, - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=stt_settings, trace_include_sensitive_data=False, @@ -679,7 +692,7 @@ async def test_session_error_event_before_session_created(): audio_input = await StreamedAudioInputFactory.get(count=2) session = OpenAISTTTranscriptionSession( input=audio_input, - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=STTModelSettings(), trace_include_sensitive_data=False, @@ -722,7 +735,7 @@ async def messages_then_timeout() -> AsyncGenerator[str, None]: audio_input = await StreamedAudioInputFactory.get(count=2) session = OpenAISTTTranscriptionSession( input=audio_input, - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=STTModelSettings(), trace_include_sensitive_data=False, @@ -778,7 +791,7 @@ async def test_inactivity_timeout(): session = OpenAISTTTranscriptionSession( input=audio_input, - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=stt_settings, trace_include_sensitive_data=False, @@ -804,7 +817,7 @@ async def test_stream_audio_buffers_turn_audio_only_for_audio_tracing( ) -> None: session = OpenAISTTTranscriptionSession( input=StreamedAudioInput(), - client=AsyncMock(api_key="FAKE_KEY"), + client=create_mock_openai_client(), model="whisper-1", settings=STTModelSettings(), trace_include_sensitive_data=False, diff --git a/tests/voice/test_openai_stt_api_key_refresh.py b/tests/voice/test_openai_stt_api_key_refresh.py new file mode 100644 index 0000000000..ae159a8d6d --- /dev/null +++ b/tests/voice/test_openai_stt_api_key_refresh.py @@ -0,0 +1,66 @@ +from typing import Any, cast +from unittest.mock import AsyncMock + +import httpx2 +import pytest +from openai import AsyncOpenAI + +from agents.voice import STTModelSettings, StreamedAudioInput +from agents.voice.models import openai_stt +from agents.voice.models.openai_stt import OpenAISTTTranscriptionSession + + +class _RotatingClient: + def __init__(self) -> None: + self.api_key = "" + self.refresh_calls = 0 + self.websocket_base_url = None + self.base_url = httpx2.URL("https://api.openai.com/v1/") + self.default_query: dict[str, str] = {} + self.auth_headers = {"Authorization": "Bearer stale"} + self.default_headers: dict[str, str] = {} + + async def _refresh_api_key(self) -> None: + self.refresh_calls += 1 + self.api_key = "sk-refreshed" + self.auth_headers = {"Authorization": f"Bearer {self.api_key}"} + + +class _WebSocketContext: + async def __aenter__(self) -> Any: + return object() + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> bool: + return False + + +@pytest.mark.asyncio +async def test_streamed_stt_refreshes_callable_api_key_before_handshake(monkeypatch) -> None: + client = _RotatingClient() + session = OpenAISTTTranscriptionSession( + input=StreamedAudioInput(), + client=cast(AsyncOpenAI, client), + model="gpt-4o-mini-transcribe", + settings=STTModelSettings(), + trace_include_sensitive_data=False, + trace_include_sensitive_audio_data=False, + ) + + captured_headers: dict[str, str] = {} + + def connect(_url: str, *, additional_headers: dict[str, str]) -> _WebSocketContext: + captured_headers.update(additional_headers) + return _WebSocketContext() + + monkeypatch.setattr(openai_stt.websockets, "connect", connect) + monkeypatch.setattr( + session, + "_setup_connection", + AsyncMock(side_effect=RuntimeError("stop after handshake")), + ) + + with pytest.raises(RuntimeError, match="stop after handshake"): + await session._process_websocket_connection() + + assert client.refresh_calls == 1 + assert captured_headers["Authorization"] == "Bearer sk-refreshed" diff --git a/tests/voice/test_openai_stt_client_config.py b/tests/voice/test_openai_stt_client_config.py new file mode 100644 index 0000000000..6813444f3f --- /dev/null +++ b/tests/voice/test_openai_stt_client_config.py @@ -0,0 +1,128 @@ +from typing import cast +from unittest.mock import MagicMock + +import httpx2 +from openai import NOT_GIVEN, AsyncOpenAI, omit + +from agents.voice.models.openai_stt import ( + _prepare_websocket_headers, + _prepare_websocket_url, +) + + +def _mock_client(**attributes: object) -> AsyncOpenAI: + attributes.setdefault("default_query", {}) + return cast(AsyncOpenAI, MagicMock(**attributes)) + + +def test_streaming_stt_websocket_url_uses_client_base_url() -> None: + client = _mock_client( + websocket_base_url=None, + base_url=httpx2.URL("https://voice-proxy.example.test/v1/"), + ) + + url = httpx2.URL(_prepare_websocket_url(client)) + + assert url.scheme == "wss" + assert url.host == "voice-proxy.example.test" + assert url.path == "/v1/realtime" + assert url.params["intent"] == "transcription" + + +def test_streaming_stt_websocket_url_prefers_websocket_base_url() -> None: + client = _mock_client( + websocket_base_url="https://voice-ws.example.test/custom/?tenant=one", + base_url=httpx2.URL("https://ignored.example.test/v1/"), + ) + + url = httpx2.URL(_prepare_websocket_url(client)) + + assert url.scheme == "wss" + assert url.host == "voice-ws.example.test" + assert url.path == "/custom/realtime" + assert url.params["tenant"] == "one" + assert url.params["intent"] == "transcription" + + +def test_streaming_stt_websocket_url_merges_client_default_query() -> None: + client = _mock_client( + websocket_base_url="wss://voice-ws.example.test/custom/?tenant=one&remove=base", + base_url=httpx2.URL("https://ignored.example.test/v1/"), + default_query={ + "api-version": "2026-08-01-preview", + "remove": omit, + "skip": NOT_GIVEN, + }, + ) + + url = httpx2.URL(_prepare_websocket_url(client)) + + assert url.params["tenant"] == "one" + assert url.params["api-version"] == "2026-08-01-preview" + assert url.params["intent"] == "transcription" + assert "remove" not in url.params + assert "skip" not in url.params + + +def test_streaming_stt_websocket_headers_use_client_configuration() -> None: + client = _mock_client( + auth_headers={"Authorization": "Bearer sk-client"}, + default_headers={ + "OpenAI-Organization": "org-client", + "OpenAI-Project": "proj-client", + "X-Proxy-Token": "proxy-token", + }, + ) + + headers = _prepare_websocket_headers(client) + + assert headers["Authorization"] == "Bearer sk-client" + assert headers["OpenAI-Organization"] == "org-client" + assert headers["OpenAI-Project"] == "proj-client" + assert headers["X-Proxy-Token"] == "proxy-token" + assert headers["OpenAI-Log-Session"] == "1" + + +def test_streaming_stt_websocket_headers_skip_openai_omission_sentinels() -> None: + client = _mock_client( + auth_headers={"Authorization": "Bearer sk-client"}, + default_headers={ + "OpenAI-Organization": omit, + "OpenAI-Project": NOT_GIVEN, + "X-Proxy-Token": "proxy-token", + }, + ) + + headers = _prepare_websocket_headers(client) + + assert headers["Authorization"] == "Bearer sk-client" + assert headers["X-Proxy-Token"] == "proxy-token" + assert "OpenAI-Organization" not in headers + assert "OpenAI-Project" not in headers + assert headers["OpenAI-Log-Session"] == "1" + + +def test_streaming_stt_websocket_headers_omit_removes_inherited_header() -> None: + client = _mock_client( + auth_headers={"Authorization": "Bearer sk-client"}, + default_headers={"authorization": omit}, + ) + + headers = _prepare_websocket_headers(client) + + assert all(key.lower() != "authorization" for key in headers) + assert headers["OpenAI-Log-Session"] == "1" + + +def test_streaming_stt_websocket_fixed_session_header_replaces_client_casing() -> None: + client = _mock_client( + auth_headers={}, + default_headers={"openai-log-session": "0"}, + ) + + headers = _prepare_websocket_headers(client) + + session_headers = { + key: value for key, value in headers.items() if key.lower() == "openai-log-session" + } + assert session_headers == {"OpenAI-Log-Session": "1"}