Skip to content
This repository was archived by the owner on Sep 1, 2026. It is now read-only.

fix: share local server across client copies #357

Open
abhinavkr26104 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: share local server across client copies #357

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 35 additions & 2 deletions src/stagehand/_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
### </END CUSTOM CODE>

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
Expand Down Expand Up @@ -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,
Expand All @@ -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

### </END CUSTOM CODE>

# Alias for `copy` for nicer inline usage, e.g.
Expand Down Expand Up @@ -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
### </END CUSTOM CODE>

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
Expand Down Expand Up @@ -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,
Expand All @@ -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

### </END CUSTOM CODE>

# Alias for `copy` for nicer inline usage, e.g.
Expand Down
41 changes: 39 additions & 2 deletions src/stagehand/_custom/sea_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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

@cubic-dev-ai cubic-dev-ai Bot Aug 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3: Every local-mode copy()/with_options() constructs a fresh SeaServerManager (plus binary-path resolution) inside self.__class__(...) before reuse_local_mode_server overwrites target._sea_server with the source's server, so that newly built manager is allocated and immediately discarded whenever reuse applies. Because the server starts lazily, this is not a process leak, but you can avoid the redundant construction by deciding reuse before building the copy (e.g., reuse the source server into the copy without letting the constructor allocate a fresh manager), or by having reuse_local_mode_server short-circuit so the constructor path for shared copies skips creating a new manager.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At src/stagehand/_custom/sea_server.py, line 409:

<comment>Every local-mode `copy()`/`with_options()` constructs a fresh `SeaServerManager` (plus binary-path resolution) inside `self.__class__(...)` before `reuse_local_mode_server` overwrites `target._sea_server` with the source's server, so that newly built manager is allocated and immediately discarded whenever reuse applies. Because the server starts lazily, this is not a process leak, but you can avoid the redundant construction by deciding reuse before building the copy (e.g., reuse the source server into the copy without letting the constructor allocate a fresh manager), or by having `reuse_local_mode_server` short-circuit so the constructor path for shared copies skips creating a new manager.</comment>

<file context>
@@ -373,6 +376,40 @@ def copy_local_mode_kwargs(
+        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
+
</file context>

target._owns_sea_server = False

@cubic-dev-ai cubic-dev-ai Bot Aug 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: If the owner client closes before a derived client, the derived client can restart the shared SEA process but can never close it because ownership stays false. Add shared reference tracking (or equivalent ownership handoff) so one active client can always shut down a restarted shared manager.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At src/stagehand/_custom/sea_server.py, line 410:

<comment>If the owner client closes before a derived client, the derived client can restart the shared SEA process but can never close it because ownership stays false. Add shared reference tracking (or equivalent ownership handoff) so one active client can always shut down a restarted shared manager.</comment>

<file context>
@@ -373,6 +376,40 @@ def copy_local_mode_kwargs(
+        and all(value is None for value in local_overrides)
+    ):
+        target._sea_server = source._sea_server
+        target._owns_sea_server = False
+
+
</file context>



def prepare_sync_client_base_url(client: _HasLocalModeState) -> str | None:
if client._sea_server is None:
return None
Expand All @@ -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()
47 changes: 47 additions & 0 deletions tests/test_local_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down