[py] Allow MAX_WS_MESSAGE_SIZE to be configurable

py/private/cdp.py

@@ -21,6 +21,7 @@
 import itertools
 import json
 import logging
+import os
 import pathlib
 from collections import defaultdict
 from collections.abc import AsyncGenerator, AsyncIterator, Generator
@@ -36,6 +37,17 @@
 T = TypeVar("T")
 MAX_WS_MESSAGE_SIZE = 2**24
 
+
+def _resolve_max_message_size(explicit=None):
+    """Return the WebSocket max message size to use.
+
+    Priority: explicit argument > ``SE_CDP_MAX_WS_MESSAGE_SIZE`` env var > ``MAX_WS_MESSAGE_SIZE``.
+    """
+    if explicit is not None:
+        return explicit
+    return int(os.environ.get("SE_CDP_MAX_WS_MESSAGE_SIZE", MAX_WS_MESSAGE_SIZE))
+
+
 devtools = None
 version = None
 
@@ -482,25 +494,30 @@ async def _reader_task(self):
 
 
 @asynccontextmanager
-async def open_cdp(url) -> AsyncIterator[CdpConnection]:
+async def open_cdp(url, max_message_size=None) -> AsyncIterator[CdpConnection]:
     """Async context manager opens a connection to the browser then closes the connection when the block exits.
 
     The context manager also sets the connection as the default
     connection for the current task, so that commands like ``await
     target.get_targets()`` will run on this connection automatically. If
     you want to use multiple connections concurrently, it is recommended
     to open each on in a separate task.
+
+    Args:
+        url: WebSocket URL of the browser's CDP endpoint.
+        max_message_size: Maximum WebSocket message size in bytes. Defaults to the
+            ``SE_CDP_MAX_WS_MESSAGE_SIZE`` environment variable, or 16 MiB if unset.
     """
     async with trio.open_nursery() as nursery:
-        conn = await connect_cdp(nursery, url)
+        conn = await connect_cdp(nursery, url, max_message_size=max_message_size)
         try:
             with connection_context(conn):
                 yield conn
         finally:
             await conn.aclose()
 
 
-async def connect_cdp(nursery, url) -> CdpConnection:
+async def connect_cdp(nursery, url, max_message_size=None) -> CdpConnection:
     """Connect to the browser specified by ``url`` and spawn a background task in the specified nursery.
 
     The ``open_cdp()`` context manager is preferred in most situations.
@@ -512,8 +529,14 @@ async def connect_cdp(nursery, url) -> CdpConnection:
     connection will be installed as the default connection for the
     current task. This argument is for unusual use cases, such as
     running inside of a notebook.
+
+    Args:
+        nursery: Trio nursery to spawn the reader task in.
+        url: WebSocket URL of the browser's CDP endpoint.
+        max_message_size: Maximum WebSocket message size in bytes. Defaults to the
+            ``SE_CDP_MAX_WS_MESSAGE_SIZE`` environment variable, or 16 MiB if unset.
     """
-    ws = await connect_websocket_url(nursery, url, max_message_size=MAX_WS_MESSAGE_SIZE)
+    ws = await connect_websocket_url(nursery, url, max_message_size=_resolve_max_message_size(max_message_size))
     cdp_conn = CdpConnection(ws)
     nursery.start_soon(cdp_conn._reader_task)
     return cdp_conn

py/selenium/webdriver/remote/client_config.py

@@ -74,6 +74,13 @@ class ClientConfig:
     """Gets and Sets the WebSocket response wait timeout (in seconds) used for communicating with the browser."""
     websocket_interval = _ClientConfigDescriptor("_websocket_interval")
     """Gets and Sets the WebSocket response wait interval (in seconds) used for communicating with the browser."""
+    websocket_max_message_size = _ClientConfigDescriptor("_websocket_max_message_size")
+    """Gets and Sets the maximum WebSocket message size in bytes for CDP connections.
+
+    Only applies to CDP-based connections (``driver.bidi_connection()``). When ``None``
+    the value falls back to the ``SE_CDP_MAX_WS_MESSAGE_SIZE`` environment variable,
+    then to the built-in default of 16 MiB (``2**24``).
+    """
 
     def __init__(
         self,
@@ -92,6 +99,7 @@ def __init__(
         extra_headers: dict | None = None,
         websocket_timeout: float | None = 30.0,
         websocket_interval: float | None = 0.1,
+        websocket_max_message_size: int | None = None,
     ) -> None:
         self.remote_server_addr = remote_server_addr
         self.keep_alive = keep_alive
@@ -107,6 +115,7 @@ def __init__(
         self.extra_headers = extra_headers
         self.websocket_timeout = websocket_timeout
         self.websocket_interval = websocket_interval
+        self.websocket_max_message_size = websocket_max_message_size
 
         self.ca_certs = (
             (os.getenv("REQUESTS_CA_BUNDLE") if "REQUESTS_CA_BUNDLE" in os.environ else certifi.where())

py/selenium/webdriver/remote/webdriver.py

@@ -1152,7 +1152,10 @@ async def bidi_connection(self):
             raise WebDriverException("Unable to find url to connect to from capabilities")
 
         devtools = cdp.import_devtools(version)
-        async with cdp.open_cdp(ws_url) as conn:
+        max_message_size = None
+        if isinstance(self.command_executor, RemoteConnection):
+            max_message_size = self.command_executor.client_config.websocket_max_message_size
+        async with cdp.open_cdp(ws_url, max_message_size=max_message_size) as conn:
             targets = await conn.execute(devtools.target.get_targets())
             for target in targets:
                 if target.target_id == self.current_window_handle:

py/test/unit/selenium/webdriver/common/cdp_ws_message_size_tests.py

@@ -0,0 +1,38 @@
+# Licensed to the Software Freedom Conservancy (SFC) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The SFC licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from selenium.webdriver.common.bidi.cdp import MAX_WS_MESSAGE_SIZE, _resolve_max_message_size
+
+
+def test_default_when_no_env_var_and_no_explicit(monkeypatch):
+    monkeypatch.delenv("SE_CDP_MAX_WS_MESSAGE_SIZE", raising=False)
+    assert _resolve_max_message_size() == MAX_WS_MESSAGE_SIZE
+
+
+def test_env_var_overrides_default(monkeypatch):
+    monkeypatch.setenv("SE_CDP_MAX_WS_MESSAGE_SIZE", str(2**26))
+    assert _resolve_max_message_size() == 2**26
+
+
+def test_explicit_param_overrides_env_var(monkeypatch):
+    monkeypatch.setenv("SE_CDP_MAX_WS_MESSAGE_SIZE", str(2**26))
+    assert _resolve_max_message_size(2**28) == 2**28
+
+
+def test_explicit_param_used_when_no_env_var(monkeypatch):
+    monkeypatch.delenv("SE_CDP_MAX_WS_MESSAGE_SIZE", raising=False)
+    assert _resolve_max_message_size(2**30) == 2**30

py/test/unit/selenium/webdriver/remote/client_config_tests.py

@@ -0,0 +1,39 @@
+# Licensed to the Software Freedom Conservancy (SFC) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The SFC licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from selenium.webdriver.remote.client_config import ClientConfig
+
+
+@pytest.fixture
+def config():
+    return ClientConfig(remote_server_addr="http://localhost:4444")
+
+
+def test_websocket_max_message_size_defaults_to_none(config):
+    assert config.websocket_max_message_size is None
+
+
+def test_websocket_max_message_size_can_be_set(config):
+    config.websocket_max_message_size = 2**26
+    assert config.websocket_max_message_size == 2**26
+
+
+def test_websocket_max_message_size_via_constructor():
+    cfg = ClientConfig(remote_server_addr="http://localhost:4444", websocket_max_message_size=2**26)
+    assert cfg.websocket_max_message_size == 2**26