diff --git a/src/openai/lib/responses_websocket/README.md b/src/openai/lib/responses_websocket/README.md index 52833cca43..9c2610d07b 100644 --- a/src/openai/lib/responses_websocket/README.md +++ b/src/openai/lib/responses_websocket/README.md @@ -186,7 +186,7 @@ it. After reconnect releases a closed default lane, `session.lane()` (or `session.lane(None)`) registers a fresh default and updates `session.default`. Limits apply only to this opt-in session. Event queues are bounded per lane and -across the session; overflow reports `ResponsesWebSocketBufferError` and closes +across the session. By default, overflow reports `ResponsesWebSocketBufferError` and closes the owned connection so ordered events are not silently dropped. Queue byte counts measure received message bytes before parsing, not Python object overhead or transport decoding memory. Response accumulation is uncapped by default: @@ -194,6 +194,24 @@ omit `max_response_bytes` or set it to `None`. A positive integer opts into a budget for the cumulative compatible event bytes passed to an accumulator; exceeding it disables that response's collection, while raw events and other lanes remain usable. Queue and lane limits still require positive integers. + +To keep healthy lanes usable when another lane falls behind, pass +`isolate_lane_overflow=True` to either session constructor. The default is `False`. +An overflowing lane loses its queued events and response collection; its reads +and final-response collection raise the same first `ResponsesWebSocketBufferError`. +Later events for that failed ID are dropped, and its reservation lasts until +reconnect. Manual lane detachment and unregistered-stream inspection keep their +existing default-lane behavior. + +For aggregate queue pressure, the session fails the largest existing backlog +in the pressured unit (bytes first, otherwise events), excluding the incoming +event. Registration order breaks ties. An individually oversized event fails +its target lane. The default lane can also fail under this policy: subsequent +default events, including API error/control events, are then unavailable, while +registered named lanes remain usable. Physical connection failures and explicit +session close still end the session. After reconnect, register a new default +explicitly with `session.lane()` if the old default failed. + Choose budgets for large legitimate outputs. These are application budgets, not protocol limits. The examples disable the transport's separate 1 MiB message limit with diff --git a/src/openai/lib/responses_websocket/_session.py b/src/openai/lib/responses_websocket/_session.py index 91dcb1297f..f7a080dece 100644 --- a/src/openai/lib/responses_websocket/_session.py +++ b/src/openai/lib/responses_websocket/_session.py @@ -157,6 +157,7 @@ def __init__(self, stream_id: str | None, limit: int | None) -> None: self.queue: deque[tuple[ResponsesServerEvent, int]] = deque() self.size = 0 self.closed = False + self.failure: ResponsesWebSocketBufferError | None = None self.in_flight = False self.closed_connection: object | None = None self.received_generation = 0 @@ -167,10 +168,15 @@ def __init__(self, stream_id: str | None, limit: int | None) -> None: class _Router: def __init__( - self, limits: ResponsesWebSocketLimits, connection: ResponsesConnection | AsyncResponsesConnection + self, + limits: ResponsesWebSocketLimits, + connection: ResponsesConnection | AsyncResponsesConnection, + *, + isolate_lane_overflow: bool = False, ) -> None: self.connection = connection self.limits = limits + self.isolate_lane_overflow = isolate_lane_overflow self.condition = threading.Condition() self.lanes: dict[str | None, _LaneState] = {} self.events = 0 @@ -213,6 +219,11 @@ def detach(self, lane: _LaneState) -> None: lane.wake.set() self.condition.notify_all() + def _fail(self, lane: _LaneState) -> None: + if not lane.closed: + lane.failure = ResponsesWebSocketBufferError("The lane exceeded its buffered event limits") + self.detach(lane) + def reconnected(self, connection: object) -> None: with self.condition: for stream_id, lane in list(self.lanes.items()): @@ -227,11 +238,21 @@ def dispatch(self, event: ResponsesServerEvent, size: int) -> None: return stream_id = _field(event, "stream_id") lane = self.lanes.get(stream_id, self.default) + if lane.failure is not None: + return if lane.closed: lane = self.default if lane.closed: return - if ( + if self.isolate_lane_overflow: + if ( + len(lane.queue) >= self.limits.max_events_per_lane + or lane.size + size > self.limits.max_bytes_per_lane + or size > self.limits.max_bytes + ): + self._fail(lane) + return + elif ( len(lane.queue) >= self.limits.max_events_per_lane or self.events >= self.limits.max_events or lane.size + size > self.limits.max_bytes_per_lane @@ -245,6 +266,18 @@ def dispatch(self, event: ResponsesServerEvent, size: int) -> None: lane.size += size self.size += size self.events += 1 + if self.isolate_lane_overflow: + while self.events > self.limits.max_events or self.size > self.limits.max_bytes: + byte_pressure = self.size > self.limits.max_bytes + largest, largest_size = lane, 0 + for candidate in self.lanes.values(): + # Exclude this incoming event once; ties follow registration order. + backlog = candidate.size if byte_pressure else len(candidate.queue) + if candidate is lane: + backlog -= size if byte_pressure else 1 + if backlog > largest_size: + largest, largest_size = candidate, backlog + self._fail(largest) lane.wake.set() self.condition.notify_all() @@ -259,6 +292,8 @@ def finish(self, error: BaseException | None = None) -> None: def pop(self, lane: _LaneState) -> ResponsesServerEvent | None: # The caller holds the condition. No await occurs after dequeueing. + if lane.failure is not None: + raise lane.failure if lane.closed: raise EOFError("The WebSocket lane is detached") if lane.queue: @@ -353,11 +388,17 @@ def get_final_response(self, *, timeout: float | None = None) -> Response: remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) self.recv(timeout=remaining) while self._state.accumulator.final is None: - if self._state.accumulator.error is not None: - raise self._state.accumulator.error + with self._session._router.condition: + if self._state.failure is not None: + raise self._state.failure + if self._state.accumulator.error is not None: + raise self._state.accumulator.error remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) self.recv(timeout=remaining) - return self._state.accumulator.final + with self._session._router.condition: + if self._state.failure is not None: + raise self._state.failure + return self._state.accumulator.final def close(self) -> None: self._session._router.detach(self._state) @@ -375,13 +416,19 @@ class ResponsesWebSocketSession: Keep a lane open for sequential responses. Reconnect releases closed lanes; open lanes remain registered on the new socket. + By default queue overflow closes the connection. ``isolate_lane_overflow=True`` + instead fails the affected lane and drops its later events. A failed default + lane also loses later default/API controls; physical failures remain global. + Recovery remains the existing connection's opt-in policy. Reconnecting a transport does not restore server-side response state or replay lane sends. """ - def __init__(self, connection: ResponsesConnection, *, limits: ResponsesWebSocketLimits) -> None: + def __init__( + self, connection: ResponsesConnection, *, limits: ResponsesWebSocketLimits, isolate_lane_overflow: bool = False + ) -> None: self.connection = connection - self._router = _Router(limits, connection) + self._router = _Router(limits, connection, isolate_lane_overflow=isolate_lane_overflow) self._thread: threading.Thread | None = None self._closed = False self._owns_connection = False @@ -515,10 +562,16 @@ async def get_final_response(self) -> Response: ): await self.recv() while self._state.accumulator.final is None: - if self._state.accumulator.error is not None: - raise self._state.accumulator.error + with self._session._router.condition: + if self._state.failure is not None: + raise self._state.failure + if self._state.accumulator.error is not None: + raise self._state.accumulator.error await self.recv() - return self._state.accumulator.final + with self._session._router.condition: + if self._state.failure is not None: + raise self._state.failure + return self._state.accumulator.final def close(self) -> None: self._session._router.detach(self._state) @@ -527,9 +580,15 @@ def close(self) -> None: class AsyncResponsesWebSocketSession: """Async counterpart of ResponsesWebSocketSession, owning the connection lifetime.""" - def __init__(self, connection: AsyncResponsesConnection, *, limits: ResponsesWebSocketLimits) -> None: + def __init__( + self, + connection: AsyncResponsesConnection, + *, + limits: ResponsesWebSocketLimits, + isolate_lane_overflow: bool = False, + ) -> None: self.connection = connection - self._router = _Router(limits, connection) + self._router = _Router(limits, connection, isolate_lane_overflow=isolate_lane_overflow) self._task: asyncio.Task[None] | None = None self._close_task: asyncio.Task[None] | None = None self._closed = False diff --git a/tests/lib/responses/test_websocket_session.py b/tests/lib/responses/test_websocket_session.py index 81e6f30344..cce9978e48 100644 --- a/tests/lib/responses/test_websocket_session.py +++ b/tests/lib/responses/test_websocket_session.py @@ -90,6 +90,7 @@ async def session_for( connection_options: dict[str, Any] | None = None, client_headers: dict[str, str] | None = None, connection_headers: Headers | None = None, + session_options: dict[str, Any] | None = None, ) -> AsyncGenerator[Driver]: options = connection_options or {} headers = {"X-Contract-Test": "synthetic", **(connection_headers or {})} @@ -102,7 +103,7 @@ async def session_for( http_client=httpx2.Client(trust_env=False), ) as client, client.responses.connect(extra_query={"contract": "1"}, extra_headers=headers, **options) as connection, - ResponsesWebSocketSession(connection, limits=limits) as session, + ResponsesWebSocketSession(connection, limits=limits, **(session_options or {})) as session, ): yield Driver(session, mode) else: @@ -116,7 +117,7 @@ async def session_for( async_client.responses.connect( extra_query={"contract": "1"}, extra_headers=headers, **options ) as async_connection, - AsyncResponsesWebSocketSession(async_connection, limits=limits) as async_session, + AsyncResponsesWebSocketSession(async_connection, limits=limits, **(session_options or {})) as async_session, ): yield Driver(async_session, mode) diff --git a/tests/lib/responses/test_websocket_session_overflow.py b/tests/lib/responses/test_websocket_session_overflow.py new file mode 100644 index 0000000000..c38242e788 --- /dev/null +++ b/tests/lib/responses/test_websocket_session_overflow.py @@ -0,0 +1,326 @@ +from __future__ import annotations + +import json +import asyncio +import threading +from dataclasses import replace + +import pytest +from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError +from websockets.sync.server import ServerConnection + +from openai.lib.responses_websocket import ResponsesWebSocketBufferError +from openai.types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides + +from .test_websocket_session import LIMITS, session_for, script_server, response_event + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize("limit", ["max_events_per_lane", "max_events", "max_bytes_per_lane", "max_bytes"]) +async def test_opt_in_overflow_fails_slow_lane_and_preserves_healthy_final(mode: str, limit: str) -> None: + slow = {"type": "response.future", "stream_id": "slow", "padding": "x" * 1000} + final = response_event("completed", "healthy") + cap = 1 if "events" in limit else len(json.dumps(slow).encode()) + 1 + + def script(socket: ServerConnection) -> None: + try: + socket.recv(timeout=5) + except ConnectionClosedOK: + return + socket.send(json.dumps(slow)) + socket.send(json.dumps(slow)) + socket.send(json.dumps(final)) + + with script_server(script) as url: + async with session_for( + mode, url, replace(LIMITS, **{limit: cap}), session_options={"isolate_lane_overflow": True} + ) as driver: + slow_lane = driver.session.lane("slow") + healthy = driver.session.lane("healthy") + await driver.call(healthy, "send", {"type": "response.create"}) + assert (await driver.call(healthy, "get_final_response")).id == "resp_healthy" + with pytest.raises(ResponsesWebSocketBufferError) as first: + await driver.call(slow_lane, "recv") + slow_lane.close() + with pytest.raises(ResponsesWebSocketBufferError) as again: + await driver.call(slow_lane, "get_final_response") + assert again.value is first.value + with pytest.raises(ValueError, match="already registered"): + driver.session.lane("slow") + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize("byte_pressure", [False, True]) +@pytest.mark.parametrize("tie", [False, True]) +async def test_aggregate_pressure_uses_existing_backlog_and_registration_ties( + mode: str, byte_pressure: bool, tie: bool +) -> None: + def event(lane: str, value: int) -> dict[str, object]: + return {"type": "response.future", "stream_id": lane, "value": value, "extra": {"kept": True}} + + frames = ( + [event("one", 1), event("two", 2), event("new", 3)] + if tie + else [event("one", 1), event("two", 2), event("two", 3), event("one", 4)] + ) + capacity = 2 if tie else 3 + limits = replace( + LIMITS, + **( + {"max_bytes": capacity * len(json.dumps(frames[0]).encode())} if byte_pressure else {"max_events": capacity} + ), + ) + + def script(socket: ServerConnection) -> None: + socket.recv(timeout=5) + for frame in frames: + socket.send(json.dumps(frame)) + socket.recv(timeout=5) + socket.send(json.dumps(event("new", 9))) + + with script_server(script) as url: + async with session_for(mode, url, limits, session_options={"isolate_lane_overflow": True}) as driver: + one = driver.session.lane("one") + two = driver.session.lane("two") + new = driver.session.lane("new") + await driver.call(new, "send", {"type": "response.create"}) + # Keep all backlogs unconsumed until the reader applies pressure. + # This is only a scheduling barrier; assertions use public lanes. + for _ in range(500): + with driver.session._router.condition: + if one._state.closed or two._state.closed: + break + await asyncio.sleep(0.01) + with pytest.raises(ResponsesWebSocketBufferError): + await driver.call(one if tie else two, "recv") + expected = [(two, 2), (new, 3)] if tie else [(one, 1), (one, 4)] + for lane, value in expected: + observed = await driver.call(lane, "recv") + assert observed.value == value + assert observed.extra == {"kept": True} + # Freed aggregate capacity admits later traffic on the same socket. + await driver.call( + new, "send", {"type": "response.steer", "previous_response_id": "synthetic", "input": "synthetic"} + ) + assert (await driver.call(new, "recv")).value == 9 + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize("limit", ["max_bytes_per_lane", "max_bytes"]) +async def test_oversized_incoming_wakes_failed_lane_without_evicting_other_backlog(mode: str, limit: str) -> None: + def script(socket: ServerConnection) -> None: + socket.recv(timeout=5) + socket.send(json.dumps({"type": "response.future", "stream_id": "other", "value": 7})) + socket.send(json.dumps({"type": "response.future", "stream_id": "large", "padding": "x" * 2000})) + socket.send(json.dumps(response_event("completed", "healthy"))) + + with script_server(script) as url: + async with session_for( + mode, url, replace(LIMITS, **{limit: 1000}), session_options={"isolate_lane_overflow": True} + ) as driver: + large = driver.session.lane("large") + other = driver.session.lane("other") + healthy = driver.session.lane("healthy") + pending = asyncio.create_task(driver.call(large, "recv")) + await asyncio.sleep(0) + await driver.call(healthy, "send", {"type": "response.create"}) + with pytest.raises(ResponsesWebSocketBufferError): + await pending + assert (await driver.call(healthy, "get_final_response")).id == "resp_healthy" + assert (await driver.call(other, "recv")).value == 7 + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_failed_id_drops_before_fallback_but_manual_detach_remains_visible(mode: str) -> None: + def script(socket: ServerConnection) -> None: + socket.recv(timeout=5) + for _ in range(8): + socket.send(json.dumps({"type": "response.future", "stream_id": "failed"})) + socket.send(json.dumps({"type": "response.future", "stream_id": "manual", "value": 4})) + socket.send(json.dumps(response_event("completed", "healthy"))) + socket.recv(timeout=5) + socket.send(json.dumps({"type": "response.future", "value": 9})) + + with script_server(script) as url: + async with session_for( + mode, url, replace(LIMITS, max_events_per_lane=1), session_options={"isolate_lane_overflow": True} + ) as driver: + failed = driver.session.lane("failed") + driver.session.lane("manual").close() + healthy = driver.session.lane("healthy") + await driver.call(healthy, "send", {"type": "response.create"}) + await driver.call(healthy, "get_final_response") + with pytest.raises(ResponsesWebSocketBufferError): + await driver.call(failed, "recv") + observed = await driver.call(driver.session.default, "recv") + assert observed.stream_id == "manual" and observed.value == 4 + await driver.call(driver.session.default, "send", {"type": "response.create"}) + assert (await driver.call(driver.session.default, "recv")).value == 9 + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_default_failure_loses_default_controls_but_transport_failure_remains_global(mode: str) -> None: + def script(socket: ServerConnection) -> None: + socket.recv(timeout=5) + for _ in range(2): + socket.send(json.dumps({"type": "response.future"})) + socket.send(json.dumps({"type": "error", "error": {"type": "synthetic", "message": "control"}})) + socket.send(json.dumps(response_event("completed", "healthy"))) + socket.close(code=1011, reason="synthetic transport failure") + + with script_server(script) as url: + async with session_for( + mode, url, replace(LIMITS, max_events_per_lane=1), session_options={"isolate_lane_overflow": True} + ) as driver: + healthy = driver.session.lane("healthy") + await driver.call(healthy, "send", {"type": "response.create"}) + assert (await driver.call(healthy, "get_final_response")).id == "resp_healthy" + with pytest.raises(ResponsesWebSocketBufferError) as first: + await driver.call(driver.session.default, "recv") + with pytest.raises(ResponsesWebSocketBufferError) as second: + await driver.call(driver.session.default, "get_final_response") + assert second.value is first.value + with pytest.raises(ConnectionClosedError): + await driver.call(healthy, "recv") + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_explicit_false_keeps_fatal_overflow_policy(mode: str) -> None: + closed = threading.Event() + + def script(socket: ServerConnection) -> None: + socket.recv(timeout=5) + for _ in range(2): + socket.send(json.dumps({"type": "response.future"})) + with pytest.raises(ConnectionClosedOK): + socket.recv(timeout=5) + closed.set() + + with script_server(script) as url: + async with session_for( + mode, url, replace(LIMITS, max_events_per_lane=1), session_options={"isolate_lane_overflow": False} + ) as driver: + await driver.call(driver.session.default, "send", {"type": "response.create"}) + assert await asyncio.to_thread(closed.wait, 5) + with pytest.raises(ResponsesWebSocketBufferError): + while True: + await driver.call(driver.session.default, "recv") + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize("default_failed", [False, True]) +async def test_failed_reservation_releases_only_after_reconnect_and_default_restores_explicitly( + mode: str, default_failed: bool +) -> None: + disconnect = threading.Event() + failed_id = None if default_failed else "worker" + observer_id = "worker" if default_failed else None + + def reconnect(_event: ReconnectingEvent) -> ReconnectingOverrides: + return {"extra_headers": {"X-Overflow-Recovery": "fresh"}} + + def script(socket: ServerConnection) -> None: + assert socket.request is not None + if socket.request.headers.get("X-Overflow-Recovery") != "fresh": + socket.recv(timeout=5) + for _ in range(2): + socket.send(json.dumps({"type": "response.future", **({"stream_id": failed_id} if failed_id else {})})) + socket.send(json.dumps(response_event("completed", observer_id))) + assert disconnect.wait(5) + socket.close(code=1011) + return + socket.send( + json.dumps( + {"type": "response.future", "restored": True, **({"stream_id": observer_id} if observer_id else {})} + ) + ) + command = json.loads(socket.recv(timeout=5)) + assert command.get("stream_id") == failed_id + socket.send(json.dumps(response_event("completed", failed_id, id="resp_restored"))) + + with script_server(script, expected_connections=2) as url: + async with session_for( + mode, + url, + replace(LIMITS, max_lanes=2, max_events_per_lane=1), + connection_options={"on_reconnecting": reconnect, "initial_delay": 0, "max_retries": 1}, + session_options={"isolate_lane_overflow": True}, + ) as driver: + worker = driver.session.lane("worker") + failed = driver.session.default if default_failed else worker + observer = worker if default_failed else driver.session.default + await driver.call(worker, "send", {"type": "response.create"}) + await driver.call(observer, "get_final_response") + with pytest.raises(ResponsesWebSocketBufferError) as first: + await driver.call(failed, "recv") + with pytest.raises(ValueError, match="already registered"): + driver.session.lane(failed_id) + disconnect.set() + assert (await driver.call(observer, "recv")).restored + new = driver.session.lane(failed_id) + if default_failed: + assert driver.session.default is new + await driver.call(new, "send", {"type": "response.create"}) + assert (await driver.call(new, "get_final_response")).id == "resp_restored" + with pytest.raises(ResponsesWebSocketBufferError) as again: + await driver.call(failed, "get_final_response") + assert again.value is first.value + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_overflow_invalidates_cached_final_with_first_buffer_error(mode: str) -> None: + def script(socket: ServerConnection) -> None: + socket.recv(timeout=5) + socket.send(json.dumps(response_event("completed", "cached"))) + socket.recv(timeout=5) + for _ in range(2): + socket.send(json.dumps({"type": "response.future", "stream_id": "cached"})) + socket.send(json.dumps(response_event("completed", "healthy"))) + + with script_server(script) as url: + async with session_for( + mode, url, replace(LIMITS, max_events_per_lane=1), session_options={"isolate_lane_overflow": True} + ) as driver: + cached = driver.session.lane("cached") + healthy = driver.session.lane("healthy") + await driver.call(cached, "send", {"type": "response.create"}) + assert (await driver.call(cached, "get_final_response")).id == "resp_cached" + await driver.call(healthy, "send", {"type": "response.create"}) + await driver.call(healthy, "get_final_response") + with pytest.raises(ResponsesWebSocketBufferError) as first: + await driver.call(cached, "get_final_response") + with pytest.raises(ResponsesWebSocketBufferError) as second: + await driver.call(cached, "recv") + assert second.value is first.value + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_queue_failure_replaces_prior_response_collection_error(mode: str) -> None: + def script(socket: ServerConnection) -> None: + socket.recv(timeout=5) + socket.send(json.dumps(response_event("created", "large"))) + socket.recv(timeout=5) + for _ in range(2): + socket.send(json.dumps({"type": "response.future", "stream_id": "large"})) + socket.send(json.dumps({"type": "response.future", "stream_id": "healthy", "value": 7})) + + with script_server(script) as url: + async with session_for( + mode, + url, + replace(LIMITS, max_events_per_lane=1, max_response_bytes=1), + session_options={"isolate_lane_overflow": True}, + ) as driver: + large = driver.session.lane("large") + healthy = driver.session.lane("healthy") + await driver.call(large, "send", {"type": "response.create"}) + assert (await driver.call(large, "recv")).type == "response.created" + with pytest.raises(ResponsesWebSocketBufferError, match="max_response_bytes"): + await driver.call(large, "get_final_response") + await driver.call(healthy, "send", {"type": "response.create"}) + assert (await driver.call(healthy, "recv")).value == 7 + with pytest.raises(ResponsesWebSocketBufferError, match="buffered event") as first: + await driver.call(large, "get_final_response") + with pytest.raises(ResponsesWebSocketBufferError) as again: + await driver.call(large, "recv") + assert again.value is first.value