feat: expose Responses WebSocket keepalive options (#3080)

Author: seratch

src/agents/__init__.py

@@ -83,7 +83,11 @@
 from .models.openai_agent_registration import OpenAIAgentRegistrationConfig
 from .models.openai_chatcompletions import OpenAIChatCompletionsModel
 from .models.openai_provider import OpenAIProvider
-from .models.openai_responses import OpenAIResponsesModel, OpenAIResponsesWSModel
+from .models.openai_responses import (
+    OpenAIResponsesModel,
+    OpenAIResponsesWebSocketOptions,
+    OpenAIResponsesWSModel,
+)
 from .prompts import DynamicPromptFunction, GenerateDynamicPromptData, Prompt
 from .repl import run_demo_loop
 from .responses_websocket_session import ResponsesWebSocketSession, responses_websocket_session
@@ -527,6 +531,7 @@ def enable_verbose_stdout_logging():
     "set_default_openai_client",
     "set_default_openai_api",
     "set_default_openai_responses_transport",
+    "OpenAIResponsesWebSocketOptions",
     "set_default_openai_harness",
     "set_default_openai_agent_registration",
     "responses_websocket_session",

src/agents/models/multi_provider.py

@@ -8,6 +8,7 @@
 from .interface import Model, ModelProvider
 from .openai_agent_registration import OpenAIAgentRegistrationConfig
 from .openai_provider import OpenAIProvider
+from .openai_responses import OpenAIResponsesWebSocketOptions
 
 MultiProviderOpenAIPrefixMode = Literal["alias", "model_id"]
 MultiProviderUnknownPrefixMode = Literal["error", "model_id"]
@@ -86,6 +87,7 @@ def __init__(
         openai_prefix_mode: MultiProviderOpenAIPrefixMode = "alias",
         unknown_prefix_mode: MultiProviderUnknownPrefixMode = "error",
         openai_agent_registration: OpenAIAgentRegistrationConfig | None = None,
+        openai_responses_websocket_options: OpenAIResponsesWebSocketOptions | None = None,
     ) -> None:
         """Create a new OpenAI provider.
 
@@ -117,6 +119,8 @@ def __init__(
                 such as ``openrouter/openai/gpt-4o``.
             openai_agent_registration: Optional agent registration configuration for the OpenAI
                 provider.
+            openai_responses_websocket_options: Optional low-level websocket keepalive options for
+                the OpenAI Responses websocket transport.
         """
         self.provider_map = provider_map
         self.openai_provider = OpenAIProvider(
@@ -129,6 +133,7 @@ def __init__(
             use_responses=openai_use_responses,
             use_responses_websocket=openai_use_responses_websocket,
             agent_registration=openai_agent_registration,
+            responses_websocket_options=openai_responses_websocket_options,
         )
         self._openai_prefix_mode = self._validate_openai_prefix_mode(openai_prefix_mode)
         self._unknown_prefix_mode = self._validate_unknown_prefix_mode(unknown_prefix_mode)

src/agents/models/openai_provider.py

@@ -16,7 +16,11 @@
     resolve_openai_agent_registration_config,
 )
 from .openai_chatcompletions import OpenAIChatCompletionsModel
-from .openai_responses import OpenAIResponsesModel, OpenAIResponsesWSModel
+from .openai_responses import (
+    OpenAIResponsesModel,
+    OpenAIResponsesWebSocketOptions,
+    OpenAIResponsesWSModel,
+)
 
 # This is kept for backward compatibility but using get_default_model() method is recommended.
 DEFAULT_MODEL: str = "gpt-4o"
@@ -49,6 +53,7 @@ def __init__(
         use_responses: bool | None = None,
         use_responses_websocket: bool | None = None,
         agent_registration: OpenAIAgentRegistrationConfig | None = None,
+        responses_websocket_options: OpenAIResponsesWebSocketOptions | None = None,
     ) -> None:
         """Create a new OpenAI provider.
 
@@ -67,6 +72,8 @@ def __init__(
             use_responses_websocket: Whether to use websocket transport for the OpenAI responses
                 API.
             agent_registration: Optional agent registration configuration.
+            responses_websocket_options: Optional low-level websocket keepalive options for the
+                OpenAI Responses websocket transport.
         """
         if openai_client is not None:
             assert api_key is None and base_url is None and websocket_base_url is None, (
@@ -95,6 +102,7 @@ def __init__(
             self._responses_transport = _openai_shared.get_default_openai_responses_transport()
         # Backward-compatibility shim for internal tests/diagnostics that inspect the legacy flag.
         self._use_responses_websocket = self._responses_transport == "websocket"
+        self._responses_websocket_options = responses_websocket_options
 
         # Reuse websocket model wrappers so websocket transport can keep a persistent connection
         # when callers pass model names as strings through a shared provider.
@@ -214,17 +222,22 @@ def get_model(self, model_name: str | None) -> Model:
         if not self._use_responses:
             return OpenAIChatCompletionsModel(model=resolved_model_name, openai_client=client)
 
-        responses_model_type = (
-            OpenAIResponsesWSModel if use_websocket_transport else OpenAIResponsesModel
-        )
-        model = responses_model_type(
+        if use_websocket_transport:
+            model = OpenAIResponsesWSModel(
+                model=resolved_model_name,
+                openai_client=client,
+                model_is_explicit=model_is_explicit,
+                websocket_options=self._responses_websocket_options,
+            )
+            if loop_cache is not None:
+                loop_cache[cache_key] = model
+            return model
+
+        model = OpenAIResponsesModel(
             model=resolved_model_name,
             openai_client=client,
             model_is_explicit=model_is_explicit,
         )
-        if use_websocket_transport:
-            if loop_cache is not None:
-                loop_cache[cache_key] = model
         return model
 
     async def aclose(self) -> None:

src/agents/models/openai_responses.py

@@ -30,6 +30,7 @@
 )
 from openai.types.responses.response_prompt_param import ResponsePromptParam
 from openai.types.responses.tool_param import LocalShell
+from typing_extensions import NotRequired
 
 from .. import _debug
 from .._tool_identity import (
@@ -191,6 +192,24 @@ class _WebsocketRequestTimeouts:
     recv: float | None
 
 
+class OpenAIResponsesWebSocketOptions(TypedDict):
+    """Low-level OpenAI Responses websocket connection options."""
+
+    ping_interval: NotRequired[float | None]
+    """Time in seconds between keepalive pings sent by the client.
+
+    The underlying ``websockets`` library usually defaults to 20.0. Set to ``None`` to
+    disable keepalive pings.
+    """
+
+    ping_timeout: NotRequired[float | None]
+    """Time in seconds to wait for a pong response before disconnecting.
+
+    Set to ``None`` to keep pings enabled but disable heartbeat timeouts during large latency
+    spikes.
+    """
+
+
 class _ResponseStreamWithRequestId:
     """Wrap an SDK event stream and retain the originating request ID."""
 
@@ -911,10 +930,14 @@ def __init__(
         openai_client: AsyncOpenAI,
         *,
         model_is_explicit: bool = True,
+        websocket_options: OpenAIResponsesWebSocketOptions | None = None,
     ) -> None:
         super().__init__(
             model=model, openai_client=openai_client, model_is_explicit=model_is_explicit
         )
+        self._websocket_options = cast(
+            OpenAIResponsesWebSocketOptions, dict(websocket_options or {})
+        )
         self._ws_connection: Any | None = None
         self._ws_connection_identity: tuple[str, tuple[tuple[str, str], ...]] | None = None
         self._ws_connection_loop_ref: weakref.ReferenceType[asyncio.AbstractEventLoop] | None = None
@@ -1531,12 +1554,20 @@ async def _open_websocket_connection(
                 "Install `websockets` or `openai[realtime]`."
             ) from exc
 
+        connect_kwargs: dict[str, Any] = {
+            "user_agent_header": None,
+            "additional_headers": dict(headers),
+            "max_size": None,
+            "open_timeout": connect_timeout,
+        }
+        if "ping_interval" in self._websocket_options:
+            connect_kwargs["ping_interval"] = self._websocket_options["ping_interval"]
+        if "ping_timeout" in self._websocket_options:
+            connect_kwargs["ping_timeout"] = self._websocket_options["ping_timeout"]
+
         return await connect(
             ws_url,
-            user_agent_header=None,
-            additional_headers=dict(headers),
-            max_size=None,
-            open_timeout=connect_timeout,
+            **connect_kwargs,
         )
 
 

src/agents/responses_websocket_session.py

@@ -13,6 +13,7 @@
     MultiProviderUnknownPrefixMode,
 )
 from .models.openai_provider import OpenAIProvider
+from .models.openai_responses import OpenAIResponsesWebSocketOptions
 from .result import RunResult, RunResultStreaming
 from .run import Runner
 from .run_config import RunConfig
@@ -86,6 +87,7 @@ async def responses_websocket_session(
     project: str | None = None,
     openai_prefix_mode: MultiProviderOpenAIPrefixMode = "alias",
     unknown_prefix_mode: MultiProviderUnknownPrefixMode = "error",
+    responses_websocket_options: OpenAIResponsesWebSocketOptions | None = None,
 ) -> AsyncIterator[ResponsesWebSocketSession]:
     """Create a shared OpenAI Responses websocket session for multiple Runner calls.
 
@@ -99,6 +101,9 @@ async def responses_websocket_session(
     configured OpenAI-compatible endpoint expects literal namespaced model IDs instead of the SDK's
     historical routing-prefix behavior.
 
+    Pass ``responses_websocket_options`` to customize low-level websocket keepalive behavior such
+    as ``ping_interval`` and ``ping_timeout``.
+
     Drain or close streamed iterators before the context exits. Exiting the context while a
     websocket request is still in flight may force-close the shared connection.
     """
@@ -112,6 +117,7 @@ async def responses_websocket_session(
         openai_use_responses_websocket=True,
         openai_prefix_mode=openai_prefix_mode,
         unknown_prefix_mode=unknown_prefix_mode,
+        openai_responses_websocket_options=responses_websocket_options,
     )
     provider = model_provider.openai_provider
     session = ResponsesWebSocketSession(

tests/test_config.py

@@ -7,6 +7,7 @@
 import pytest
 
 from agents import (
+    responses_websocket_session,
     set_default_openai_api,
     set_default_openai_client,
     set_default_openai_key,
@@ -170,6 +171,36 @@ async def test_openai_provider_reuses_websocket_model_instance_for_same_model_na
     assert model1 is model2
 
 
+@pytest.mark.asyncio
+async def test_openai_provider_passes_responses_websocket_options_to_model():
+    class DummyAsyncOpenAI:
+        pass
+
+    provider = OpenAIProvider(
+        use_responses=True,
+        use_responses_websocket=True,
+        openai_client=DummyAsyncOpenAI(),  # type: ignore[arg-type]
+        responses_websocket_options={"ping_interval": 30.0, "ping_timeout": None},
+    )
+
+    model = provider.get_model("gpt-4")
+
+    assert isinstance(model, OpenAIResponsesWSModel)
+    assert model._websocket_options == {"ping_interval": 30.0, "ping_timeout": None}
+
+
+@pytest.mark.asyncio
+async def test_responses_websocket_session_passes_keepalive_options_to_provider():
+    async with responses_websocket_session(
+        api_key="test-key",
+        responses_websocket_options={"ping_interval": None, "ping_timeout": None},
+    ) as session:
+        assert session.provider._responses_websocket_options == {
+            "ping_interval": None,
+            "ping_timeout": None,
+        }
+
+
 def test_openai_provider_does_not_reuse_non_websocket_model_instances():
     provider = OpenAIProvider(use_responses=True, use_responses_websocket=False)
 

tests/test_openai_responses.py

@@ -1580,6 +1580,40 @@ async def fake_open(
     assert ws.sent_messages[1]["previous_response_id"] == "resp-1"
 
 
+@pytest.mark.asyncio
+async def test_websocket_model_passes_keepalive_options_to_connect(monkeypatch):
+    import websockets.asyncio.client as websockets_client
+
+    client = DummyWSClient()
+    model = OpenAIResponsesWSModel(
+        model="gpt-4",
+        openai_client=client,  # type: ignore[arg-type]
+        websocket_options={"ping_interval": 45.0, "ping_timeout": None},
+    )
+    ws = DummyWSConnection([])
+    captured_kwargs: dict[str, Any] = {}
+
+    async def fake_connect(ws_url: str, **kwargs: Any) -> DummyWSConnection:
+        captured_kwargs["ws_url"] = ws_url
+        captured_kwargs.update(kwargs)
+        return ws
+
+    monkeypatch.setattr(websockets_client, "connect", fake_connect)
+
+    opened = await model._open_websocket_connection(
+        "wss://example.test/v1/responses",
+        {"Authorization": "Bearer test-key"},
+        connect_timeout=10.0,
+    )
+
+    assert opened is ws
+    assert captured_kwargs["ws_url"] == "wss://example.test/v1/responses"
+    assert captured_kwargs["additional_headers"] == {"Authorization": "Bearer test-key"}
+    assert captured_kwargs["open_timeout"] == 10.0
+    assert captured_kwargs["ping_interval"] == 45.0
+    assert captured_kwargs["ping_timeout"] is None
+
+
 @pytest.mark.allow_call_model_methods
 def test_websocket_model_reconnects_when_reused_from_different_event_loop(monkeypatch):
     client = DummyWSClient()