diff --git a/src/openai/_exceptions.py b/src/openai/_exceptions.py index 7a30e4a336..d54e4942fb 100644 --- a/src/openai/_exceptions.py +++ b/src/openai/_exceptions.py @@ -1,7 +1,9 @@ from __future__ import annotations +import sys +from datetime import timedelta from typing import TYPE_CHECKING, Any, Optional, cast -from typing_extensions import Literal +from typing_extensions import Literal, override import httpx2 @@ -30,9 +32,205 @@ "WebSocketQueueFullError", ] +_RequestTimeoutSnapshot = dict[str, float | None] +_RequestSnapshot = tuple[str, list[tuple[bytes, bytes]], _RequestTimeoutSnapshot | None] +_ResponseExtensionsSnapshot = dict[str, bytes] +_ResponseSnapshot = tuple[ + int, + list[tuple[bytes, bytes]], + _RequestSnapshot | None, + timedelta | None, + _ResponseExtensionsSnapshot, +] + +_REDACTED_REQUEST_URL = "https://redacted.invalid/" +_REDACTED_HEADER_VALUE = b"" +_SAFE_HEADER_VALUES = { + b"content-length", + b"content-type", + b"x-request-id", +} +_TIMEOUT_EXTENSION_KEYS = ("connect", "read", "write", "pool") +_RESPONSE_EXTENSION_KEYS = ("http_version", "reason_phrase") + class OpenAIError(Exception): - pass + @override + def __reduce__(self) -> tuple[object, tuple[object, ...]]: + state = self.__dict__.copy() + slot_state = _snapshot_slot_state(self) + + if isinstance(self, WebSocketConnectionClosedError): + state["unsent_messages"] = [] + + request_snapshot: _RequestSnapshot | None = None + request = state.get("request") + if _is_http_request(request): + request_snapshot = _snapshot_request(request) + del state["request"] + + response_snapshot: _ResponseSnapshot | None = None + response = state.get("response") + if _is_http_response(response): + response_snapshot = _snapshot_response(response) + del state["response"] + + return ( + _reconstruct_openai_error, + (type(self), self.args, state, slot_state, request_snapshot, response_snapshot), + ) + + +def _is_legacy_httpx_instance(value: object, type_name: str) -> bool: + module = sys.modules.get("httpx") + legacy_type = getattr(module, type_name, None) if module is not None else None + return isinstance(legacy_type, type) and isinstance(value, legacy_type) + + +def _is_http_request(value: object) -> bool: + return isinstance(value, httpx2.Request) or _is_legacy_httpx_instance(value, "Request") + + +def _is_http_response(value: object) -> bool: + return isinstance(value, httpx2.Response) or _is_legacy_httpx_instance(value, "Response") + + +def _snapshot_headers(headers: Any) -> list[tuple[bytes, bytes]]: + return [ + (name, value if name.lower() in _SAFE_HEADER_VALUES else _REDACTED_HEADER_VALUE) + for name, value in headers.raw + ] + + +def _snapshot_timeout_extension(request: Any) -> _RequestTimeoutSnapshot | None: + extensions = request.extensions + timeout = extensions.get("timeout") if isinstance(extensions, dict) else None + if not isinstance(timeout, dict): + return None + + snapshot: _RequestTimeoutSnapshot = {} + for key in _TIMEOUT_EXTENSION_KEYS: + value = timeout.get(key) + if value is None: + if key in timeout: + snapshot[key] = None + elif isinstance(value, (int, float)): + snapshot[key] = float(value) + return snapshot or None + + +def _snapshot_request(request: Any) -> _RequestSnapshot: + return ( + request.method, + _snapshot_headers(request.headers), + _snapshot_timeout_extension(request), + ) + + +def _restore_request(snapshot: _RequestSnapshot) -> httpx2.Request: + method, headers, timeout = snapshot + kwargs: dict[str, Any] = {"headers": headers} + if timeout is not None: + kwargs["extensions"] = {"timeout": timeout} + return httpx2.Request(method, _REDACTED_REQUEST_URL, **kwargs) + + +def _snapshot_response_extensions(response: Any) -> _ResponseExtensionsSnapshot: + extensions = response.extensions + if not isinstance(extensions, dict): + return {} + + snapshot: _ResponseExtensionsSnapshot = {} + for key in _RESPONSE_EXTENSION_KEYS: + value = extensions.get(key) + if isinstance(value, bytes): + snapshot[key] = value + return snapshot + + +def _snapshot_response(response: Any) -> _ResponseSnapshot: + try: + request_snapshot = _snapshot_request(response.request) + except RuntimeError: + request_snapshot = None + + try: + elapsed = response.elapsed + except RuntimeError: + elapsed = None + + return ( + response.status_code, + _snapshot_headers(response.headers), + request_snapshot, + elapsed, + _snapshot_response_extensions(response), + ) + + +def _restore_response( + snapshot: _ResponseSnapshot, + *, + request: httpx2.Request | None, +) -> httpx2.Response: + status_code, headers, response_request_snapshot, elapsed, extensions = snapshot + if request is None and response_request_snapshot is not None: + request = _restore_request(response_request_snapshot) + + kwargs: dict[str, Any] = {"headers": headers} + if request is not None: + kwargs["request"] = request + if extensions: + kwargs["extensions"] = extensions + + response = httpx2.Response(status_code, **kwargs) + if elapsed is not None: + response.elapsed = elapsed + return response + + +def _snapshot_slot_state(error: OpenAIError) -> dict[str, Any]: + slot_state: dict[str, Any] = {} + for cls in type(error).__mro__: + slots = cls.__dict__.get("__slots__") + if slots is None: + continue + if isinstance(slots, str): + slots = (slots,) + for slot in slots: + if slot in {"__dict__", "__weakref__"}: + continue + name = slot + if name.startswith("__") and not name.endswith("__"): + name = f"_{cls.__name__.lstrip('_')}{name}" + try: + slot_state[name] = getattr(error, name) + except AttributeError: + pass + return slot_state + + +def _reconstruct_openai_error( + error_type: type[OpenAIError], + args: tuple[object, ...], + state: dict[str, Any], + slot_state: dict[str, Any], + request_snapshot: _RequestSnapshot | None, + response_snapshot: _ResponseSnapshot | None, +) -> OpenAIError: + error = Exception.__new__(error_type) + Exception.__init__(error, *args) + error.__dict__.update(state) + for name, value in slot_state.items(): + setattr(error, name, value) + + request = _restore_request(request_snapshot) if request_snapshot is not None else None + if request is not None: + error.__dict__["request"] = request + if response_snapshot is not None: + error.__dict__["response"] = _restore_response(response_snapshot, request=request) + + return error class SubjectTokenProviderError(OpenAIError): diff --git a/tests/test_exceptions_pickle.py b/tests/test_exceptions_pickle.py new file mode 100644 index 0000000000..7f85c688d1 --- /dev/null +++ b/tests/test_exceptions_pickle.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +import asyncio +import pickle +from collections.abc import Iterator +from datetime import timedelta +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Thread + +import httpx2 +import pytest + +from openai import APITimeoutError, AsyncOpenAI, BadRequestError, OpenAI, OpenAIError +from openai._exceptions import SubjectTokenProviderError, WebSocketConnectionClosedError + + +_ERROR_PAYLOAD = b'{"error":{"message":"bad request","type":"invalid_request_error","param":null,"code":"bad_request"}}' +_REDACTED_REQUEST_URL = "https://redacted.invalid/" + + +class _CustomNewOpenAIError(OpenAIError): + def __new__(cls, message: str) -> _CustomNewOpenAIError: + return Exception.__new__(cls) + + +class _SlottedOpenAIError(OpenAIError): + __slots__ = ("context",) + + def __init__(self, message: str, context: str) -> None: + super().__init__(message, context) + self.context = context + + +class _BadRequestHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self) -> None: + self.send_response(400) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(_ERROR_PAYLOAD))) + self.send_header("x-request-id", "req_transport") + self.end_headers() + self.wfile.write(_ERROR_PAYLOAD) + + def log_message(self, format: str, *args: object) -> None: + pass + + +@pytest.fixture +def transport_error_base_url() -> Iterator[str]: + server = ThreadingHTTPServer(("127.0.0.1", 0), _BadRequestHandler) + thread = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}/v1" + finally: + server.shutdown() + thread.join() + server.server_close() + + +def test_timeout_error_pickle_round_trip() -> None: + request = httpx2.Request("POST", "https://example.test/v1/responses") + error = APITimeoutError(request) + + restored = pickle.loads(pickle.dumps(error)) + + assert isinstance(restored, APITimeoutError) + assert str(restored) == "Request timed out." + assert restored.message == "Request timed out." + assert restored.request.method == "POST" + assert str(restored.request.url) == _REDACTED_REQUEST_URL + assert restored.body is None + + +def test_request_private_data_is_omitted_and_timeout_is_preserved() -> None: + request = httpx2.Request( + "POST", + "https://example.test/v1/responses?customer_secret=query-private", + content=b"request-body-private", + extensions={ + "timeout": { + "connect": 1.0, + "read": 2.0, + "write": 3.0, + "pool": 4.0, + } + }, + ) + error = APITimeoutError(request) + + payload = pickle.dumps(error) + restored = pickle.loads(payload) + + assert b"query-private" not in payload + assert b"request-body-private" not in payload + assert str(restored.request.url) == _REDACTED_REQUEST_URL + assert restored.request.extensions["timeout"] == { + "connect": 1.0, + "read": 2.0, + "write": 3.0, + "pool": 4.0, + } + + +@pytest.mark.parametrize( + ("header_name", "header_value"), + [ + ("Authorization", "Bearer sk-openai-secret"), + ("api-key", "azure-secret-key"), + ("OpenAI-Organization", "org-private-identifier"), + ("OpenAI-Project", "proj-private-identifier"), + ("X-Customer-Metadata", "arbitrary-customer-data"), + ], +) +def test_request_private_header_values_are_redacted(header_name: str, header_value: str) -> None: + request = httpx2.Request( + "POST", + "https://example.test/v1/responses", + headers={header_name: header_value}, + ) + error = APITimeoutError(request) + + payload = pickle.dumps(error) + restored = pickle.loads(payload) + + assert header_value.encode() not in payload + assert restored.request.headers[header_name] == "" + + +def test_response_private_data_is_omitted_and_headers_are_redacted() -> None: + request = httpx2.Request("POST", "https://example.test/v1/token") + response = httpx2.Response( + 400, + request=request, + headers={ + "Authorization": "Bearer response-auth-secret", + "api-key": "response-api-key-secret", + "Set-Cookie": "session=response-cookie-secret", + "x-request-id": "req_123", + }, + content=b"response-body-private", + ) + error = SubjectTokenProviderError("token exchange failed", response=response) + + payload = pickle.dumps(error) + restored = pickle.loads(payload) + + assert b"response-auth-secret" not in payload + assert b"response-api-key-secret" not in payload + assert b"response-cookie-secret" not in payload + assert b"response-body-private" not in payload + assert restored.response is not None + assert restored.response.headers["Authorization"] == "" + assert restored.response.headers["api-key"] == "" + assert restored.response.headers["Set-Cookie"] == "" + assert restored.response.headers["x-request-id"] == "req_123" + assert restored.response.content == b"" + + +def test_response_protocol_metadata_is_preserved() -> None: + request = httpx2.Request("GET", "https://example.test/v1/models") + response = httpx2.Response( + 418, + request=request, + extensions={ + "http_version": b"HTTP/2", + "reason_phrase": b"Custom reason", + "network_stream": object(), + }, + ) + error = SubjectTokenProviderError("request failed", response=response) + + payload = pickle.dumps(error) + restored = pickle.loads(payload) + + assert restored.response is not None + assert restored.response.http_version == "HTTP/2" + assert restored.response.reason_phrase == "Custom reason" + assert "network_stream" not in restored.response.extensions + + +def test_status_error_pickle_round_trip_preserves_response_state() -> None: + request = httpx2.Request("POST", "https://example.test/v1/responses") + response = httpx2.Response( + 400, + request=request, + headers={"x-request-id": "req_123"}, + json={"error": "bad request"}, + ) + response.elapsed = timedelta(milliseconds=125) + body = { + "code": "invalid_value", + "param": "input", + "type": "invalid_request_error", + } + error = BadRequestError("Bad request", response=response, body=body) + + restored = pickle.loads(pickle.dumps(error)) + + assert isinstance(restored, BadRequestError) + assert restored.status_code == 400 + assert restored.request_id == "req_123" + assert restored.body == body + assert restored.code == "invalid_value" + assert restored.param == "input" + assert restored.type == "invalid_request_error" + assert restored.response.status_code == 400 + assert restored.response.content == b"" + assert restored.response.elapsed == timedelta(milliseconds=125) + + +def _assert_transport_status_error_pickle_round_trip(error: BadRequestError) -> None: + restored = pickle.loads(pickle.dumps(error)) + + assert isinstance(restored, BadRequestError) + assert restored.status_code == 400 + assert restored.request_id == "req_transport" + assert restored.response.status_code == 400 + assert restored.response.request is restored.request + assert restored.request.method == "GET" + assert str(restored.request.url) == _REDACTED_REQUEST_URL + assert restored.body == { + "error": { + "message": "bad request", + "type": "invalid_request_error", + "param": None, + "code": "bad_request", + } + } + assert restored.response.content == b"" + assert restored.request.headers["Authorization"] == "" + assert error.request.headers["Authorization"] != restored.request.headers["Authorization"] + assert restored.response.elapsed >= timedelta(0) + + +def test_status_error_from_sync_transport_pickle_round_trip(transport_error_base_url: str) -> None: + with OpenAI(api_key="test", base_url=transport_error_base_url) as client: + with pytest.raises(BadRequestError) as exc_info: + client.models.list() + + _assert_transport_status_error_pickle_round_trip(exc_info.value) + + +def test_status_error_from_async_transport_pickle_round_trip(transport_error_base_url: str) -> None: + async def run() -> None: + async with AsyncOpenAI(api_key="test", base_url=transport_error_base_url) as client: + with pytest.raises(BadRequestError) as exc_info: + await client.models.list() + + _assert_transport_status_error_pickle_round_trip(exc_info.value) + + asyncio.run(run()) + + +def test_custom_error_subclass_with_required_new_argument_pickle_round_trip() -> None: + error = _CustomNewOpenAIError("custom error") + + restored = pickle.loads(pickle.dumps(error)) + + assert isinstance(restored, _CustomNewOpenAIError) + assert str(restored) == "custom error" + + +def test_custom_error_subclass_preserves_slotted_state() -> None: + error = _SlottedOpenAIError("custom error", "slot context") + + restored = pickle.loads(pickle.dumps(error)) + + assert isinstance(restored, _SlottedOpenAIError) + assert restored.args == ("custom error", "slot context") + assert restored.context == "slot context" + + +def test_websocket_error_pickle_omits_unsent_payloads() -> None: + sensitive_event = '{"type":"input_audio_buffer.append","audio":"private-audio"}' + error = WebSocketConnectionClosedError( + "connection closed", + unsent_messages=[sensitive_event], + ) + + payload = pickle.dumps(error) + restored = pickle.loads(payload) + + assert isinstance(restored, WebSocketConnectionClosedError) + assert str(restored) == "connection closed" + assert b"private-audio" not in payload + assert restored.unsent_messages == [] diff --git a/tests/test_httpx_compat.py b/tests/test_httpx_compat.py index cb00e9b00f..b283535fe1 100644 --- a/tests/test_httpx_compat.py +++ b/tests/test_httpx_compat.py @@ -1,16 +1,18 @@ from __future__ import annotations -import os import asyncio import importlib +import os +import pickle import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any, cast -from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler from typing_extensions import override +import httpx2 import pytest -from openai import OpenAI, AsyncOpenAI +from openai import AsyncOpenAI, BadRequestError, OpenAI pytestmark = pytest.mark.skipif( os.environ.get("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires the dedicated legacy HTTPX compatibility lane" @@ -53,6 +55,56 @@ async def handler(request: Any) -> Any: assert response.data is not None +def _assert_legacy_error_pickle_round_trip(error: BadRequestError) -> None: + payload = pickle.dumps(error) + restored = pickle.loads(payload) + + assert b"legacy-secret" not in payload + assert isinstance(restored.request, httpx2.Request) + assert isinstance(restored.response, httpx2.Response) + assert restored.request.headers["Authorization"] == "" + assert str(restored.request.url) == "https://redacted.invalid/" + assert restored.response.request is restored.request + assert restored.status_code == 400 + assert restored.body == {"error": "legacy bad request"} + + +def test_legacy_sync_httpx_error_pickle_round_trip() -> None: + httpx = cast(Any, importlib.import_module("httpx")) + + def handler(request: Any) -> Any: + return httpx.Response(400, request=request, json={"error": "legacy bad request"}) + + with OpenAI( + api_key="legacy-secret", + base_url="https://example.test/v1", + http_client=httpx.Client(transport=httpx.MockTransport(handler), trust_env=False), + max_retries=0, + ) as client: + with pytest.raises(BadRequestError) as exc_info: + client.models.list() + + _assert_legacy_error_pickle_round_trip(exc_info.value) + + +async def test_legacy_async_httpx_error_pickle_round_trip() -> None: + httpx = cast(Any, importlib.import_module("httpx")) + + async def handler(request: Any) -> Any: + return httpx.Response(400, request=request, json={"error": "legacy bad request"}) + + async with AsyncOpenAI( + api_key="legacy-secret", + base_url="https://example.test/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler), trust_env=False), + max_retries=0, + ) as client: + with pytest.raises(BadRequestError) as exc_info: + await client.models.list() + + _assert_legacy_error_pickle_round_trip(exc_info.value) + + async def test_external_legacy_aiohttp_client_is_supported() -> None: httpx = cast(Any, importlib.import_module("httpx")) HttpxAiohttpClient = cast(Any, importlib.import_module("httpx_aiohttp")).HttpxAiohttpClient