From 4dc6c0473b66c060c3eaa73376d7fc76b84f60ff Mon Sep 17 00:00:00 2001 From: AutomatedTester Date: Wed, 27 May 2026 13:53:03 +0100 Subject: [PATCH] [py] Allow MAX_WS_MESSAGE_SIZE to be configurable fixes #16546 --- py/private/cdp.py | 31 +++++++++++++-- py/selenium/webdriver/remote/client_config.py | 9 +++++ py/selenium/webdriver/remote/webdriver.py | 5 ++- .../common/cdp_ws_message_size_tests.py | 38 ++++++++++++++++++ .../webdriver/remote/client_config_tests.py | 39 +++++++++++++++++++ 5 files changed, 117 insertions(+), 5 deletions(-) create mode 100644 py/test/unit/selenium/webdriver/common/cdp_ws_message_size_tests.py create mode 100644 py/test/unit/selenium/webdriver/remote/client_config_tests.py diff --git a/py/private/cdp.py b/py/private/cdp.py index 86341b3babd71..b93d773d7944a 100644 --- a/py/private/cdp.py +++ b/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,7 +494,7 @@ 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 @@ -490,9 +502,14 @@ async def open_cdp(url) -> AsyncIterator[CdpConnection]: 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 @@ -500,7 +517,7 @@ async def open_cdp(url) -> AsyncIterator[CdpConnection]: 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 diff --git a/py/selenium/webdriver/remote/client_config.py b/py/selenium/webdriver/remote/client_config.py index 98019b8f15834..a4a930e2f844a 100644 --- a/py/selenium/webdriver/remote/client_config.py +++ b/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()) diff --git a/py/selenium/webdriver/remote/webdriver.py b/py/selenium/webdriver/remote/webdriver.py index c8177068efd1b..0b965bba4cdaa 100644 --- a/py/selenium/webdriver/remote/webdriver.py +++ b/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: diff --git a/py/test/unit/selenium/webdriver/common/cdp_ws_message_size_tests.py b/py/test/unit/selenium/webdriver/common/cdp_ws_message_size_tests.py new file mode 100644 index 0000000000000..5eef7d112f6a4 --- /dev/null +++ b/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 diff --git a/py/test/unit/selenium/webdriver/remote/client_config_tests.py b/py/test/unit/selenium/webdriver/remote/client_config_tests.py new file mode 100644 index 0000000000000..5a113885ea417 --- /dev/null +++ b/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