From f5e83b9020b5b225e1dbb2753568739cc112763d Mon Sep 17 00:00:00 2001 From: Edward Liang Date: Tue, 16 Jun 2026 15:09:37 -0700 Subject: [PATCH 1/2] feat: add Valkey cache configuration to client and session start Signed-off-by: Edward Liang --- src/stagehand/_client.py | 70 ++++ src/stagehand/_custom/sea_server.py | 97 +++++ src/stagehand/_custom/session.py | 4 + src/stagehand/resources/sessions.py | 8 + src/stagehand/types/__init__.py | 1 + src/stagehand/types/session_start_params.py | 34 ++ tests/api_resources/test_sessions.py | 18 + tests/test_valkey_cache.py | 440 ++++++++++++++++++++ 8 files changed, 672 insertions(+) create mode 100644 tests/test_valkey_cache.py diff --git a/src/stagehand/_client.py b/src/stagehand/_client.py index 0c405416..58f7a86a 100644 --- a/src/stagehand/_client.py +++ b/src/stagehand/_client.py @@ -92,6 +92,13 @@ class Stagehand(SyncAPIClient): _local_ready_timeout_s: float _local_shutdown_on_close: bool _sea_server: SeaServerManager | None + _valkey_host: str | None + _valkey_port: int | None + _valkey_tls: bool | None + _valkey_password: str | None + _valkey_username: str | None + _valkey_cache_ttl: int | None + _valkey_key_prefix: str | None ### ### @@ -109,6 +116,13 @@ def __init__( local_chrome_path: str | None = None, local_ready_timeout_s: float = 10.0, local_shutdown_on_close: bool = True, + valkey_host: str | None = None, + valkey_port: int | None = None, + valkey_tls: bool | None = None, + valkey_password: str | None = None, + valkey_username: str | None = None, + valkey_cache_ttl: int | None = None, + valkey_key_prefix: str | None = None, base_url: str | httpx.URL | None = None, timeout: float | Timeout | None | NotGiven = not_given, max_retries: int = DEFAULT_MAX_RETRIES, @@ -159,6 +173,13 @@ def __init__( local_shutdown_on_close=local_shutdown_on_close, base_url=base_url, model_api_key=model_api_key, + valkey_host=valkey_host, + valkey_port=valkey_port, + valkey_tls=valkey_tls, + valkey_password=valkey_password, + valkey_username=valkey_username, + valkey_cache_ttl=valkey_cache_ttl, + valkey_key_prefix=valkey_key_prefix, ) ### @@ -269,6 +290,13 @@ def copy( local_chrome_path: str | None = None, local_ready_timeout_s: float | None = None, local_shutdown_on_close: bool | None = None, + valkey_host: str | None = None, + valkey_port: int | None = None, + valkey_tls: bool | None = None, + valkey_password: str | None = None, + valkey_username: str | None = None, + valkey_cache_ttl: int | None = None, + valkey_key_prefix: str | None = None, base_url: str | httpx.URL | None = None, timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx.Client | None = None, @@ -323,6 +351,13 @@ def copy( local_chrome_path=local_chrome_path, local_ready_timeout_s=local_ready_timeout_s, local_shutdown_on_close=local_shutdown_on_close, + valkey_host=valkey_host, + valkey_port=valkey_port, + valkey_tls=valkey_tls, + valkey_password=valkey_password, + valkey_username=valkey_username, + valkey_cache_ttl=valkey_cache_ttl, + valkey_key_prefix=valkey_key_prefix, ), **_extra_kwargs, ) @@ -383,6 +418,13 @@ class AsyncStagehand(AsyncAPIClient): _local_ready_timeout_s: float _local_shutdown_on_close: bool _sea_server: SeaServerManager | None + _valkey_host: str | None + _valkey_port: int | None + _valkey_tls: bool | None + _valkey_password: str | None + _valkey_username: str | None + _valkey_cache_ttl: int | None + _valkey_key_prefix: str | None ### ### @@ -400,6 +442,13 @@ def __init__( local_chrome_path: str | None = None, local_ready_timeout_s: float = 10.0, local_shutdown_on_close: bool = True, + valkey_host: str | None = None, + valkey_port: int | None = None, + valkey_tls: bool | None = None, + valkey_password: str | None = None, + valkey_username: str | None = None, + valkey_cache_ttl: int | None = None, + valkey_key_prefix: str | None = None, base_url: str | httpx.URL | None = None, timeout: float | Timeout | None | NotGiven = not_given, max_retries: int = DEFAULT_MAX_RETRIES, @@ -450,6 +499,13 @@ def __init__( local_shutdown_on_close=local_shutdown_on_close, base_url=base_url, model_api_key=model_api_key, + valkey_host=valkey_host, + valkey_port=valkey_port, + valkey_tls=valkey_tls, + valkey_password=valkey_password, + valkey_username=valkey_username, + valkey_cache_ttl=valkey_cache_ttl, + valkey_key_prefix=valkey_key_prefix, ) ### @@ -560,6 +616,13 @@ def copy( local_chrome_path: str | None = None, local_ready_timeout_s: float | None = None, local_shutdown_on_close: bool | None = None, + valkey_host: str | None = None, + valkey_port: int | None = None, + valkey_tls: bool | None = None, + valkey_password: str | None = None, + valkey_username: str | None = None, + valkey_cache_ttl: int | None = None, + valkey_key_prefix: str | None = None, base_url: str | httpx.URL | None = None, timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx.AsyncClient | None = None, @@ -614,6 +677,13 @@ def copy( local_chrome_path=local_chrome_path, local_ready_timeout_s=local_ready_timeout_s, local_shutdown_on_close=local_shutdown_on_close, + valkey_host=valkey_host, + valkey_port=valkey_port, + valkey_tls=valkey_tls, + valkey_password=valkey_password, + valkey_username=valkey_username, + valkey_cache_ttl=valkey_cache_ttl, + valkey_key_prefix=valkey_key_prefix, ), **_extra_kwargs, ) diff --git a/src/stagehand/_custom/sea_server.py b/src/stagehand/_custom/sea_server.py index a72ed591..858d6504 100644 --- a/src/stagehand/_custom/sea_server.py +++ b/src/stagehand/_custom/sea_server.py @@ -7,6 +7,7 @@ import signal import socket import asyncio +import warnings import subprocess from pathlib import Path from threading import Lock @@ -28,6 +29,13 @@ class SeaServerConfig: model_api_key: str | None chrome_path: str | None shutdown_on_close: bool + valkey_host: str | None = None + valkey_port: int | None = None + valkey_tls: bool | None = None + valkey_password: str | None = None + valkey_username: str | None = None + valkey_cache_ttl: int | None = None + valkey_key_prefix: str | None = None class _HasLocalModeState(Protocol): @@ -40,6 +48,13 @@ class _HasLocalModeState(Protocol): _local_ready_timeout_s: float _local_shutdown_on_close: bool _sea_server: SeaServerManager | None + _valkey_host: str | None + _valkey_port: int | None + _valkey_tls: bool | None + _valkey_password: str | None + _valkey_username: str | None + _valkey_cache_ttl: int | None + _valkey_key_prefix: str | None class LocalModeKwargs(TypedDict): @@ -51,6 +66,13 @@ class LocalModeKwargs(TypedDict): local_chrome_path: str | None local_ready_timeout_s: float local_shutdown_on_close: bool + valkey_host: str | None + valkey_port: int | None + valkey_tls: bool | None + valkey_password: str | None + valkey_username: str | None + valkey_cache_ttl: int | None + valkey_key_prefix: str | None def _pick_free_port(host: str) -> int: @@ -163,6 +185,28 @@ def _build_process_env(self, *, port: int) -> dict[str, str]: if self._config.chrome_path: proc_env["CHROME_PATH"] = self._config.chrome_path proc_env["LIGHTHOUSE_CHROMIUM_PATH"] = self._config.chrome_path + if self._config.valkey_host is not None: + proc_env["VALKEY_HOST"] = self._config.valkey_host + if self._config.valkey_port is not None: + proc_env["VALKEY_PORT"] = str(self._config.valkey_port) + if self._config.valkey_tls is not None: + proc_env["VALKEY_TLS"] = "true" if self._config.valkey_tls else "false" + if self._config.valkey_password is not None: + host = self._config.valkey_host or "" + if not self._config.valkey_tls and host not in ("localhost", "127.0.0.1", "::1"): + warnings.warn( + "valkey_password is set but valkey_tls is not enabled. " + "Credentials will be sent in cleartext. Set valkey_tls=True for non-local hosts.", + UserWarning, + stacklevel=6, + ) + proc_env["VALKEY_PASSWORD"] = self._config.valkey_password + if self._config.valkey_username is not None: + proc_env["VALKEY_USERNAME"] = self._config.valkey_username + if self._config.valkey_cache_ttl is not None: + proc_env["CACHE_TTL"] = str(self._config.valkey_cache_ttl) + if self._config.valkey_key_prefix is not None: + proc_env["VALKEY_KEY_PREFIX"] = self._config.valkey_key_prefix return proc_env def ensure_running_sync(self) -> str: @@ -299,6 +343,13 @@ def configure_client_base_url( local_shutdown_on_close: bool, base_url: str | httpx.URL | None, model_api_key: str | None, + valkey_host: str | None = None, + valkey_port: int | None = None, + valkey_tls: bool | None = None, + valkey_password: str | None = None, + valkey_username: str | None = None, + valkey_cache_ttl: int | None = None, + valkey_key_prefix: str | None = None, ) -> str | httpx.URL: client._server_mode = server client._local_stagehand_binary_path = _local_stagehand_binary_path @@ -308,9 +359,34 @@ def configure_client_base_url( client._local_chrome_path = local_chrome_path client._local_ready_timeout_s = local_ready_timeout_s client._local_shutdown_on_close = local_shutdown_on_close + client._valkey_host = valkey_host + client._valkey_port = valkey_port + client._valkey_tls = valkey_tls + client._valkey_password = valkey_password + client._valkey_username = valkey_username + client._valkey_cache_ttl = valkey_cache_ttl + client._valkey_key_prefix = valkey_key_prefix client._sea_server = None + _valkey_params_given = any( + p is not None + for p in (valkey_host, valkey_port, valkey_tls, valkey_password, valkey_username, valkey_cache_ttl, valkey_key_prefix) + ) + if server != "local" and _valkey_params_given: + warnings.warn( + "Client-level Valkey parameters (valkey_host, valkey_port, etc.) are only used in " + 'server="local" mode. In remote mode, pass valkey_cache= to sessions.start() instead. ' + "These parameters will have no effect.", + UserWarning, + stacklevel=3, + ) + if server == "local": + # NOTE: Valkey caching in local mode requires the @valkey/valkey-glide native + # addon to be loadable by the SEA binary. Pre-built binaries do not bundle + # native addons, so the cache connection will silently fall back to disabled. + # Use server="remote" for Valkey caching, or supply a custom binary built + # with @valkey/valkey-glide installed alongside it. if base_url is None: base_url = "http://127.0.0.1" @@ -323,6 +399,13 @@ def configure_client_base_url( model_api_key=model_api_key, chrome_path=local_chrome_path, shutdown_on_close=local_shutdown_on_close, + valkey_host=valkey_host, + valkey_port=valkey_port, + valkey_tls=valkey_tls, + valkey_password=valkey_password, + valkey_username=valkey_username, + valkey_cache_ttl=valkey_cache_ttl, + valkey_key_prefix=valkey_key_prefix, ), _local_stagehand_binary_path=_local_stagehand_binary_path, ) @@ -346,6 +429,13 @@ def copy_local_mode_kwargs( local_chrome_path: str | None, local_ready_timeout_s: float | None, local_shutdown_on_close: bool | None, + valkey_host: str | None = None, + valkey_port: int | None = None, + valkey_tls: bool | None = None, + valkey_password: str | None = None, + valkey_username: str | None = None, + valkey_cache_ttl: int | None = None, + valkey_key_prefix: str | None = None, ) -> LocalModeKwargs: return { "server": server or client._server_mode, @@ -370,6 +460,13 @@ def copy_local_mode_kwargs( if local_shutdown_on_close is not None else client._local_shutdown_on_close ), + "valkey_host": valkey_host if valkey_host is not None else client._valkey_host, + "valkey_port": valkey_port if valkey_port is not None else client._valkey_port, + "valkey_tls": valkey_tls if valkey_tls is not None else client._valkey_tls, + "valkey_password": valkey_password if valkey_password is not None else client._valkey_password, + "valkey_username": valkey_username if valkey_username is not None else client._valkey_username, + "valkey_cache_ttl": valkey_cache_ttl if valkey_cache_ttl is not None else client._valkey_cache_ttl, + "valkey_key_prefix": valkey_key_prefix if valkey_key_prefix is not None else client._valkey_key_prefix, } diff --git a/src/stagehand/_custom/session.py b/src/stagehand/_custom/session.py index 4ca99711..51ce341c 100644 --- a/src/stagehand/_custom/session.py +++ b/src/stagehand/_custom/session.py @@ -732,6 +732,7 @@ def _sync_start( experimental: bool | Omit = omit, self_heal: bool | Omit = omit, system_prompt: str | Omit = omit, + valkey_cache: session_start_params.ValkeyCacheOptions | Omit = omit, verbose: Literal[0, 1, 2] | Omit = omit, wait_for_captcha_solves: bool | Omit = omit, x_stream_response: Literal["true", "false"] | Omit = omit, @@ -751,6 +752,7 @@ def _sync_start( experimental=experimental, self_heal=self_heal, system_prompt=system_prompt, + valkey_cache=valkey_cache, verbose=verbose, wait_for_captcha_solves=wait_for_captcha_solves, x_stream_response=x_stream_response, @@ -776,6 +778,7 @@ async def _async_start( experimental: bool | Omit = omit, self_heal: bool | Omit = omit, system_prompt: str | Omit = omit, + valkey_cache: session_start_params.ValkeyCacheOptions | Omit = omit, verbose: Literal[0, 1, 2] | Omit = omit, wait_for_captcha_solves: bool | Omit = omit, x_stream_response: Literal["true", "false"] | Omit = omit, @@ -795,6 +798,7 @@ async def _async_start( experimental=experimental, self_heal=self_heal, system_prompt=system_prompt, + valkey_cache=valkey_cache, verbose=verbose, wait_for_captcha_solves=wait_for_captcha_solves, x_stream_response=x_stream_response, diff --git a/src/stagehand/resources/sessions.py b/src/stagehand/resources/sessions.py index a3c3c4bd..d722d41d 100644 --- a/src/stagehand/resources/sessions.py +++ b/src/stagehand/resources/sessions.py @@ -929,6 +929,7 @@ def start( experimental: bool | Omit = omit, self_heal: bool | Omit = omit, system_prompt: str | Omit = omit, + valkey_cache: session_start_params.ValkeyCacheOptions | Omit = omit, verbose: Literal[0, 1, 2] | Omit = omit, wait_for_captcha_solves: bool | Omit = omit, x_stream_response: Literal["true", "false"] | Omit = omit, @@ -961,6 +962,8 @@ def start( system_prompt: Custom system prompt for AI operations + valkey_cache: Valkey cache backend configuration. When set, uses Valkey for caching. + verbose: Logging verbosity level (0=quiet, 1=normal, 2=debug) wait_for_captcha_solves: Wait for captcha solves (deprecated, v2 only) @@ -999,6 +1002,7 @@ def start( "experimental": experimental, "self_heal": self_heal, "system_prompt": system_prompt, + "valkey_cache": valkey_cache, "verbose": verbose, "wait_for_captcha_solves": wait_for_captcha_solves, }, @@ -1892,6 +1896,7 @@ async def start( experimental: bool | Omit = omit, self_heal: bool | Omit = omit, system_prompt: str | Omit = omit, + valkey_cache: session_start_params.ValkeyCacheOptions | Omit = omit, verbose: Literal[0, 1, 2] | Omit = omit, wait_for_captcha_solves: bool | Omit = omit, x_stream_response: Literal["true", "false"] | Omit = omit, @@ -1924,6 +1929,8 @@ async def start( system_prompt: Custom system prompt for AI operations + valkey_cache: Valkey cache backend configuration. When set, uses Valkey for caching. + verbose: Logging verbosity level (0=quiet, 1=normal, 2=debug) wait_for_captcha_solves: Wait for captcha solves (deprecated, v2 only) @@ -1962,6 +1969,7 @@ async def start( "experimental": experimental, "self_heal": self_heal, "system_prompt": system_prompt, + "valkey_cache": valkey_cache, "verbose": verbose, "wait_for_captcha_solves": wait_for_captcha_solves, }, diff --git a/src/stagehand/types/__init__.py b/src/stagehand/types/__init__.py index d01a70d0..a495eb06 100644 --- a/src/stagehand/types/__init__.py +++ b/src/stagehand/types/__init__.py @@ -8,6 +8,7 @@ from .session_act_response import SessionActResponse as SessionActResponse from .session_end_response import SessionEndResponse as SessionEndResponse from .session_start_params import SessionStartParams as SessionStartParams +from .session_start_params import ValkeyCacheOptions as ValkeyCacheOptions from .session_execute_params import SessionExecuteParams as SessionExecuteParams from .session_extract_params import SessionExtractParams as SessionExtractParams from .session_observe_params import SessionObserveParams as SessionObserveParams diff --git a/src/stagehand/types/session_start_params.py b/src/stagehand/types/session_start_params.py index c13fd036..a8ff7632 100644 --- a/src/stagehand/types/session_start_params.py +++ b/src/stagehand/types/session_start_params.py @@ -24,9 +24,40 @@ "BrowserbaseSessionCreateParamsProxiesProxyConfigListBrowserbaseProxyConfig", "BrowserbaseSessionCreateParamsProxiesProxyConfigListBrowserbaseProxyConfigGeolocation", "BrowserbaseSessionCreateParamsProxiesProxyConfigListExternalProxyConfig", + "ValkeyCacheOptions", ] +class ValkeyCacheOptions(TypedDict, total=False): + """Configuration for the Valkey cache backend.""" + + host: Required[Annotated[str, PropertyInfo(alias="valkeyHost")]] + """Valkey host address. Required to enable Valkey caching.""" + + port: Annotated[int, PropertyInfo(alias="valkeyPort")] + """Valkey port (default: 6379).""" + + tls: Annotated[bool, PropertyInfo(alias="valkeyTls")] + """Enable TLS for the Valkey connection. + + Defaults to True when password is set, False otherwise. Set explicitly to + False when connecting to a non-TLS Valkey instance with password auth + (e.g. local development). + """ + + password: Annotated[str, PropertyInfo(alias="valkeyPassword")] + """Valkey authentication password.""" + + username: Annotated[str, PropertyInfo(alias="valkeyUsername")] + """Valkey authentication username (for ACL-enabled instances).""" + + cache_ttl: Annotated[int, PropertyInfo(alias="cacheTtl")] + """TTL in seconds for cache entries. Omit for no expiry.""" + + key_prefix: Annotated[str, PropertyInfo(alias="valkeyKeyPrefix")] + """Key prefix namespace (default: "stagehand").""" + + class SessionStartParams(TypedDict, total=False): model_name: Required[Annotated[str, PropertyInfo(alias="modelName")]] """Model name to use for AI operations""" @@ -60,6 +91,9 @@ class SessionStartParams(TypedDict, total=False): wait_for_captcha_solves: Annotated[bool, PropertyInfo(alias="waitForCaptchaSolves")] """Wait for captcha solves (deprecated, v2 only)""" + valkey_cache: Annotated[ValkeyCacheOptions, PropertyInfo(alias="valkeyCache")] + """Valkey cache backend configuration. When set, uses Valkey for caching.""" + x_stream_response: Annotated[Literal["true", "false"], PropertyInfo(alias="x-stream-response")] """Whether to stream the response via SSE""" diff --git a/tests/api_resources/test_sessions.py b/tests/api_resources/test_sessions.py index 9b9eac2e..843264f3 100644 --- a/tests/api_resources/test_sessions.py +++ b/tests/api_resources/test_sessions.py @@ -1196,6 +1196,15 @@ def test_method_start_with_all_params(self, client: Stagehand) -> None: experimental=True, self_heal=True, system_prompt="systemPrompt", + valkey_cache={ + "host": "valkey.example.com", + "port": 6380, + "tls": True, + "password": "secret", + "username": "admin", + "cache_ttl": 86400, + "key_prefix": "myapp", + }, verbose=1, wait_for_captcha_solves=True, x_stream_response="true", @@ -2404,6 +2413,15 @@ async def test_method_start_with_all_params(self, async_client: AsyncStagehand) experimental=True, self_heal=True, system_prompt="systemPrompt", + valkey_cache={ + "host": "valkey.example.com", + "port": 6380, + "tls": True, + "password": "secret", + "username": "admin", + "cache_ttl": 86400, + "key_prefix": "myapp", + }, verbose=1, wait_for_captcha_solves=True, x_stream_response="true", diff --git a/tests/test_valkey_cache.py b/tests/test_valkey_cache.py new file mode 100644 index 00000000..eec505e8 --- /dev/null +++ b/tests/test_valkey_cache.py @@ -0,0 +1,440 @@ +"""Unit tests for Valkey cache backend configuration. + +Tests cover: +- SeaServerConfig dataclass accepts valkey fields +- _build_process_env propagates valkey config as environment variables +- Client constructors accept valkey params and pass them through +- copy()/with_options() preserves and overrides valkey config +- ValkeyCacheOptions TypedDict is importable +""" + +from __future__ import annotations + +import warnings +from pathlib import Path + +import pytest + +from stagehand import Stagehand, AsyncStagehand +from stagehand.types import ValkeyCacheOptions +from stagehand._custom.sea_server import SeaServerConfig, SeaServerManager + + +class TestSeaServerConfigValkey: + def test_defaults_to_none(self) -> None: + config = SeaServerConfig( + host="127.0.0.1", + port=9222, + headless=True, + ready_timeout_s=10.0, + model_api_key=None, + chrome_path=None, + shutdown_on_close=True, + ) + assert config.valkey_host is None + assert config.valkey_port is None + assert config.valkey_tls is None + assert config.valkey_password is None + assert config.valkey_username is None + assert config.valkey_cache_ttl is None + assert config.valkey_key_prefix is None + + def test_accepts_all_valkey_fields(self) -> None: + config = SeaServerConfig( + host="127.0.0.1", + port=9222, + headless=True, + ready_timeout_s=10.0, + model_api_key=None, + chrome_path=None, + shutdown_on_close=True, + valkey_host="valkey.example.com", + valkey_port=6380, + valkey_tls=True, + valkey_password="secret", + valkey_username="admin", + valkey_cache_ttl=3600, + valkey_key_prefix="myapp", + ) + assert config.valkey_host == "valkey.example.com" + assert config.valkey_port == 6380 + assert config.valkey_tls is True + assert config.valkey_password == "secret" + assert config.valkey_username == "admin" + assert config.valkey_cache_ttl == 3600 + assert config.valkey_key_prefix == "myapp" + + +class TestBuildProcessEnvValkey: + def _make_manager(self, tmp_path: Path, **valkey_kwargs: object) -> SeaServerManager: + binary = tmp_path / "fake-binary" + binary.write_text("x") + config = SeaServerConfig( + host="127.0.0.1", + port=9222, + headless=True, + ready_timeout_s=10.0, + model_api_key="test-key", + chrome_path=None, + shutdown_on_close=True, + **valkey_kwargs, # type: ignore[arg-type] + ) + return SeaServerManager(config=config, _local_stagehand_binary_path=binary) + + def test_no_valkey_env_when_not_configured(self, tmp_path: Path) -> None: + mgr = self._make_manager(tmp_path) + env = mgr._build_process_env(port=9222) + assert "VALKEY_HOST" not in env + assert "VALKEY_PORT" not in env + assert "VALKEY_TLS" not in env + assert "VALKEY_PASSWORD" not in env + assert "VALKEY_USERNAME" not in env + assert "CACHE_TTL" not in env + assert "VALKEY_KEY_PREFIX" not in env + + def test_valkey_host_sets_env(self, tmp_path: Path) -> None: + mgr = self._make_manager(tmp_path, valkey_host="valkey.local") + env = mgr._build_process_env(port=9222) + assert env["VALKEY_HOST"] == "valkey.local" + + def test_valkey_port_sets_env(self, tmp_path: Path) -> None: + mgr = self._make_manager(tmp_path, valkey_port=6380) + env = mgr._build_process_env(port=9222) + assert env["VALKEY_PORT"] == "6380" + + def test_valkey_tls_true(self, tmp_path: Path) -> None: + mgr = self._make_manager(tmp_path, valkey_tls=True) + env = mgr._build_process_env(port=9222) + assert env["VALKEY_TLS"] == "true" + + def test_valkey_tls_false(self, tmp_path: Path) -> None: + mgr = self._make_manager(tmp_path, valkey_tls=False) + env = mgr._build_process_env(port=9222) + assert env["VALKEY_TLS"] == "false" + + def test_valkey_credentials_set_env(self, tmp_path: Path) -> None: + mgr = self._make_manager(tmp_path, valkey_host="localhost", valkey_password="pw", valkey_username="user") + env = mgr._build_process_env(port=9222) + assert env["VALKEY_PASSWORD"] == "pw" + assert env["VALKEY_USERNAME"] == "user" + + def test_valkey_cache_ttl_sets_env(self, tmp_path: Path) -> None: + mgr = self._make_manager(tmp_path, valkey_cache_ttl=7200) + env = mgr._build_process_env(port=9222) + assert env["CACHE_TTL"] == "7200" + + def test_valkey_key_prefix_sets_env(self, tmp_path: Path) -> None: + mgr = self._make_manager(tmp_path, valkey_key_prefix="custom") + env = mgr._build_process_env(port=9222) + assert env["VALKEY_KEY_PREFIX"] == "custom" + + def test_all_valkey_env_vars_together(self, tmp_path: Path) -> None: + mgr = self._make_manager( + tmp_path, + valkey_host="host", + valkey_port=6381, + valkey_tls=True, # TLS enabled so password-without-TLS warning won't fire + valkey_password="pass", + valkey_username="usr", + valkey_cache_ttl=300, + valkey_key_prefix="pfx", + ) + env = mgr._build_process_env(port=9222) + assert env["VALKEY_HOST"] == "host" + assert env["VALKEY_PORT"] == "6381" + assert env["VALKEY_TLS"] == "true" + assert env["VALKEY_PASSWORD"] == "pass" + assert env["VALKEY_USERNAME"] == "usr" + assert env["CACHE_TTL"] == "300" + assert env["VALKEY_KEY_PREFIX"] == "pfx" + + +class TestClientValkeyParams: + def test_sync_client_accepts_valkey_params(self) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + client = Stagehand( + base_url="http://localhost:4010", + browserbase_api_key="test", + model_api_key="test", + valkey_host="valkey.local", + valkey_port=6380, + valkey_tls=True, + valkey_password="pw", + valkey_username="user", + valkey_cache_ttl=3600, + valkey_key_prefix="myapp", + ) + assert client._valkey_host == "valkey.local" + assert client._valkey_port == 6380 + assert client._valkey_tls is True + assert client._valkey_password == "pw" + assert client._valkey_username == "user" + assert client._valkey_cache_ttl == 3600 + assert client._valkey_key_prefix == "myapp" + client.close() + + async def test_async_client_accepts_valkey_params(self) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + client = AsyncStagehand( + base_url="http://localhost:4010", + browserbase_api_key="test", + model_api_key="test", + valkey_host="valkey.local", + valkey_port=6380, + ) + assert client._valkey_host == "valkey.local" + assert client._valkey_port == 6380 + await client.close() + + def test_sync_client_defaults_none(self) -> None: + client = Stagehand( + base_url="http://localhost:4010", + browserbase_api_key="test", + model_api_key="test", + ) + assert client._valkey_host is None + assert client._valkey_port is None + assert client._valkey_tls is None + assert client._valkey_password is None + assert client._valkey_username is None + assert client._valkey_cache_ttl is None + assert client._valkey_key_prefix is None + client.close() + + def test_with_options_preserves_valkey(self) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + client = Stagehand( + base_url="http://localhost:4010", + browserbase_api_key="test", + model_api_key="test", + valkey_host="original.host", + valkey_cache_ttl=1800, + ) + copied = client.with_options(max_retries=5) + assert copied._valkey_host == "original.host" + assert copied._valkey_cache_ttl == 1800 + client.close() + copied.close() + + def test_with_options_overrides_valkey(self) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + client = Stagehand( + base_url="http://localhost:4010", + browserbase_api_key="test", + model_api_key="test", + valkey_host="original.host", + valkey_cache_ttl=1800, + ) + copied = client.with_options(valkey_host="new.host", valkey_cache_ttl=900) + assert copied._valkey_host == "new.host" + assert copied._valkey_cache_ttl == 900 + client.close() + copied.close() + + def test_remote_mode_warns_when_valkey_params_set(self) -> None: + with pytest.warns(UserWarning, match="server.*local.*mode"): + client = Stagehand( + base_url="http://localhost:4010", + browserbase_api_key="test", + model_api_key="test", + valkey_host="valkey.remote.example", + ) + client.close() + + +class TestValkeyCacheOptionsType: + def test_importable(self) -> None: + # ValkeyCacheOptions should be importable from stagehand.types + assert ValkeyCacheOptions is not None + + def test_is_typed_dict(self) -> None: + # Verify it behaves as a TypedDict (has __annotations__) + assert "host" in ValkeyCacheOptions.__annotations__ + assert "port" in ValkeyCacheOptions.__annotations__ + assert "tls" in ValkeyCacheOptions.__annotations__ + assert "password" in ValkeyCacheOptions.__annotations__ + assert "username" in ValkeyCacheOptions.__annotations__ + assert "cache_ttl" in ValkeyCacheOptions.__annotations__ + assert "key_prefix" in ValkeyCacheOptions.__annotations__ + + +class TestSessionStartValkeyCache: + """Verify valkey_cache is a first-class parameter in sessions.start().""" + + def test_start_accepts_valkey_cache_param(self) -> None: + """sessions.start() should accept valkey_cache without needing extra_body.""" + import inspect + + from stagehand.resources.sessions import SessionsResource, AsyncSessionsResource + + sync_sig = inspect.signature(SessionsResource.start) + assert "valkey_cache" in sync_sig.parameters + + async_sig = inspect.signature(AsyncSessionsResource.start) + assert "valkey_cache" in async_sig.parameters + + def test_start_sends_valkey_cache_in_body(self) -> None: + """valkey_cache dict appears in the serialized request body.""" + from unittest.mock import patch + + client = Stagehand( + base_url="http://localhost:4010", + browserbase_api_key="test", + model_api_key="test", + ) + captured: list[dict[str, object]] = [] + + original_post = client.post + + def capturing_post(*args: object, **kwargs: object) -> object: + captured.append(dict(kwargs)) + return original_post(*args, **kwargs) + + with patch.object(client, "post", side_effect=capturing_post): + try: + client.sessions.start( + model_name="anthropic/claude-sonnet-4-6", + valkey_cache={"host": "valkey.local", "port": 6380}, + ) + except Exception: + pass + + assert captured, "post() was never called" + body = captured[0].get("body") + assert body is not None + assert isinstance(body, dict) + assert "valkeyCache" in body or "valkey_cache" in body or "valkeyHost" in str(body) + client.close() + + +class TestConfigureClientBaseUrlLocalMode: + """T1: configure_client_base_url local-mode branch creates SeaServerManager with valkey config.""" + + def test_local_mode_creates_sea_server_with_valkey(self, tmp_path: Path) -> None: + binary = tmp_path / "fake-binary" + binary.write_text("x") + + client = Stagehand( + base_url="http://127.0.0.1", + browserbase_api_key="test", + model_api_key="test", + server="local", + _local_stagehand_binary_path=binary, + valkey_host="valkey.local", + valkey_port=6380, + valkey_cache_ttl=300, + ) + + assert client._sea_server is not None + assert client._sea_server._config.valkey_host == "valkey.local" + assert client._sea_server._config.valkey_port == 6380 + assert client._sea_server._config.valkey_cache_ttl == 300 + client.close() + + +class TestCopyLocalModeKwargsPartialOverride: + """T2: copy_local_mode_kwargs correctly overrides one valkey param while keeping others.""" + + def test_partial_override(self) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + client = Stagehand( + base_url="http://localhost:4010", + browserbase_api_key="test", + model_api_key="test", + valkey_host="original.host", + valkey_port=6379, + valkey_cache_ttl=1800, + ) + copied = client.with_options(valkey_host="new.host") + assert copied._valkey_host == "new.host" + assert copied._valkey_port == 6379 + assert copied._valkey_cache_ttl == 1800 + client.close() + copied.close() + + +class TestFalsyButValidValues: + """T3: falsy-but-valid values (port=0, cache_ttl=0) are propagated.""" + + def test_port_zero_sets_env(self, tmp_path: Path) -> None: + binary = tmp_path / "fake-binary" + binary.write_text("x") + config = SeaServerConfig( + host="127.0.0.1", + port=9222, + headless=True, + ready_timeout_s=10.0, + model_api_key=None, + chrome_path=None, + shutdown_on_close=True, + valkey_host="localhost", + valkey_port=0, + ) + mgr = SeaServerManager(config=config, _local_stagehand_binary_path=binary) + env = mgr._build_process_env(port=9222) + assert env["VALKEY_PORT"] == "0" + + def test_cache_ttl_zero_sets_env(self, tmp_path: Path) -> None: + binary = tmp_path / "fake-binary" + binary.write_text("x") + config = SeaServerConfig( + host="127.0.0.1", + port=9222, + headless=True, + ready_timeout_s=10.0, + model_api_key=None, + chrome_path=None, + shutdown_on_close=True, + valkey_cache_ttl=0, + ) + mgr = SeaServerManager(config=config, _local_stagehand_binary_path=binary) + env = mgr._build_process_env(port=9222) + assert env["CACHE_TTL"] == "0" + + +class TestEmptyStringValues: + """T4: empty string values pass is-not-None check and produce empty env vars.""" + + def test_empty_host_sets_env(self, tmp_path: Path) -> None: + binary = tmp_path / "fake-binary" + binary.write_text("x") + config = SeaServerConfig( + host="127.0.0.1", + port=9222, + headless=True, + ready_timeout_s=10.0, + model_api_key=None, + chrome_path=None, + shutdown_on_close=True, + valkey_host="", + ) + mgr = SeaServerManager(config=config, _local_stagehand_binary_path=binary) + env = mgr._build_process_env(port=9222) + assert "VALKEY_HOST" in env + assert env["VALKEY_HOST"] == "" + + def test_empty_password_sets_env(self, tmp_path: Path) -> None: + binary = tmp_path / "fake-binary" + binary.write_text("x") + config = SeaServerConfig( + host="127.0.0.1", + port=9222, + headless=True, + ready_timeout_s=10.0, + model_api_key=None, + chrome_path=None, + shutdown_on_close=True, + valkey_host="localhost", + valkey_password="", + ) + mgr = SeaServerManager(config=config, _local_stagehand_binary_path=binary) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + env = mgr._build_process_env(port=9222) + assert "VALKEY_PASSWORD" in env + assert env["VALKEY_PASSWORD"] == "" From e4bbc4d0268fe83f2fea195621ecd32c4c624a3a Mon Sep 17 00:00:00 2001 From: Edward Liang Date: Fri, 19 Jun 2026 15:55:48 -0700 Subject: [PATCH 2/2] fix: address cubic comments Signed-off-by: Edward Liang --- src/stagehand/types/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/stagehand/types/__init__.py b/src/stagehand/types/__init__.py index a495eb06..f4cf53a5 100644 --- a/src/stagehand/types/__init__.py +++ b/src/stagehand/types/__init__.py @@ -8,7 +8,10 @@ from .session_act_response import SessionActResponse as SessionActResponse from .session_end_response import SessionEndResponse as SessionEndResponse from .session_start_params import SessionStartParams as SessionStartParams + +### from .session_start_params import ValkeyCacheOptions as ValkeyCacheOptions +### from .session_execute_params import SessionExecuteParams as SessionExecuteParams from .session_extract_params import SessionExtractParams as SessionExtractParams from .session_observe_params import SessionObserveParams as SessionObserveParams