diff --git a/src/stagehand/_client.py b/src/stagehand/_client.py index 0c405416..508a9682 100644 --- a/src/stagehand/_client.py +++ b/src/stagehand/_client.py @@ -40,6 +40,7 @@ from ._custom.session import install_stainless_session_patches from ._custom.sea_server import ( copy_local_mode_kwargs, + reuse_local_mode_server, configure_client_base_url, close_sync_client_sea_server, prepare_sync_client_base_url, @@ -92,6 +93,7 @@ class Stagehand(SyncAPIClient): _local_ready_timeout_s: float _local_shutdown_on_close: bool _sea_server: SeaServerManager | None + _owns_sea_server: bool ### ### @@ -301,7 +303,7 @@ def copy( params = set_default_query http_client = http_client or self._client - return self.__class__( + copied = self.__class__( browserbase_api_key=browserbase_api_key or self.browserbase_api_key, browserbase_project_id=browserbase_project_id or self.browserbase_project_id, model_api_key=model_api_key or self.model_api_key, @@ -326,6 +328,21 @@ def copy( ), **_extra_kwargs, ) + reuse_local_mode_server( + self, + copied, + server=server, + model_api_key=model_api_key, + _local_stagehand_binary_path=_local_stagehand_binary_path, + local_host=local_host, + local_port=local_port, + local_headless=local_headless, + local_chrome_path=local_chrome_path, + local_ready_timeout_s=local_ready_timeout_s, + local_shutdown_on_close=local_shutdown_on_close, + ) + return copied + ### # Alias for `copy` for nicer inline usage, e.g. @@ -383,6 +400,7 @@ class AsyncStagehand(AsyncAPIClient): _local_ready_timeout_s: float _local_shutdown_on_close: bool _sea_server: SeaServerManager | None + _owns_sea_server: bool ### ### @@ -592,7 +610,7 @@ def copy( params = set_default_query http_client = http_client or self._client - return self.__class__( + copied = self.__class__( browserbase_api_key=browserbase_api_key or self.browserbase_api_key, browserbase_project_id=browserbase_project_id or self.browserbase_project_id, model_api_key=model_api_key or self.model_api_key, @@ -617,6 +635,21 @@ def copy( ), **_extra_kwargs, ) + reuse_local_mode_server( + self, + copied, + server=server, + model_api_key=model_api_key, + _local_stagehand_binary_path=_local_stagehand_binary_path, + local_host=local_host, + local_port=local_port, + local_headless=local_headless, + local_chrome_path=local_chrome_path, + local_ready_timeout_s=local_ready_timeout_s, + local_shutdown_on_close=local_shutdown_on_close, + ) + return copied + ### # Alias for `copy` for nicer inline usage, e.g. diff --git a/src/stagehand/_custom/sea_server.py b/src/stagehand/_custom/sea_server.py index a72ed591..ad8ec36b 100644 --- a/src/stagehand/_custom/sea_server.py +++ b/src/stagehand/_custom/sea_server.py @@ -40,6 +40,7 @@ class _HasLocalModeState(Protocol): _local_ready_timeout_s: float _local_shutdown_on_close: bool _sea_server: SeaServerManager | None + _owns_sea_server: bool class LocalModeKwargs(TypedDict): @@ -309,6 +310,7 @@ def configure_client_base_url( client._local_ready_timeout_s = local_ready_timeout_s client._local_shutdown_on_close = local_shutdown_on_close client._sea_server = None + client._owns_sea_server = False if server == "local": if base_url is None: @@ -326,6 +328,7 @@ def configure_client_base_url( ), _local_stagehand_binary_path=_local_stagehand_binary_path, ) + client._owns_sea_server = True return base_url if base_url is None: @@ -373,6 +376,40 @@ def copy_local_mode_kwargs( } +def reuse_local_mode_server( + source: _HasLocalModeState, + target: _HasLocalModeState, + *, + server: Literal["remote", "local"] | None, + model_api_key: str | None, + _local_stagehand_binary_path: str | os.PathLike[str] | None, + local_host: str | None, + local_port: int | None, + local_headless: bool | None, + local_chrome_path: str | None, + local_ready_timeout_s: float | None, + local_shutdown_on_close: bool | None, +) -> None: + """Share an unchanged local server with a derived client.""" + local_overrides = ( + model_api_key, + _local_stagehand_binary_path, + local_host, + local_port, + local_headless, + local_chrome_path, + local_ready_timeout_s, + local_shutdown_on_close, + ) + if ( + source._sea_server is not None + and (server is None or server == "local") + and all(value is None for value in local_overrides) + ): + target._sea_server = source._sea_server + target._owns_sea_server = False + + def prepare_sync_client_base_url(client: _HasLocalModeState) -> str | None: if client._sea_server is None: return None @@ -386,10 +423,10 @@ async def prepare_async_client_base_url(client: _HasLocalModeState) -> str | Non def close_sync_client_sea_server(client: _HasLocalModeState) -> None: - if client._sea_server is not None: + if client._sea_server is not None and client._owns_sea_server: client._sea_server.close() async def close_async_client_sea_server(client: _HasLocalModeState) -> None: - if client._sea_server is not None: + if client._sea_server is not None and client._owns_sea_server: await client._sea_server.aclose() diff --git a/tests/test_local_server.py b/tests/test_local_server.py index 12af5df4..36f87cd8 100644 --- a/tests/test_local_server.py +++ b/tests/test_local_server.py @@ -54,6 +54,53 @@ def kill(self) -> None: self._returncode = 0 +def test_sync_copy_reuses_local_server_without_taking_ownership(monkeypatch: pytest.MonkeyPatch) -> None: + _set_required_env(monkeypatch) + client = Stagehand(server="local", _local_stagehand_binary_path="/does/not/matter/in/test") + dummy = _DummySeaServer("http://127.0.0.1:43123") + client._sea_server = dummy # type: ignore[assignment] + + copied = client.with_options(max_retries=3) + + assert copied._sea_server is dummy + assert copied._owns_sea_server is False + copied.close() + assert dummy.closed == 0 + client.close() + assert dummy.closed == 1 + + +@pytest.mark.asyncio +async def test_async_copy_reuses_local_server_without_taking_ownership(monkeypatch: pytest.MonkeyPatch) -> None: + _set_required_env(monkeypatch) + client = AsyncStagehand(server="local", _local_stagehand_binary_path="/does/not/matter/in/test") + dummy = _DummySeaServer("http://127.0.0.1:43123") + client._sea_server = dummy # type: ignore[assignment] + + copied = client.with_options(timeout=3) + + assert copied._sea_server is dummy + assert copied._owns_sea_server is False + await copied.close() + assert dummy.closed == 0 + await client.close() + assert dummy.closed == 1 + + +def test_copy_with_local_process_override_uses_independent_server(monkeypatch: pytest.MonkeyPatch) -> None: + _set_required_env(monkeypatch) + client = Stagehand(server="local", _local_stagehand_binary_path="/does/not/matter/in/test") + dummy = _DummySeaServer("http://127.0.0.1:43123") + client._sea_server = dummy # type: ignore[assignment] + + copied = client.with_options(local_port=43124) + + assert copied._sea_server is not dummy + assert copied._owns_sea_server is True + copied.close() + client.close() + + def _set_required_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("BROWSERBASE_API_KEY", "bb_key")