Skip to content

More custom code cleanup #334

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pirate merged 4 commits into from
Apr 3, 2026
+50 −46
Merged

More custom code cleanup #334

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5971304
remove unneeded file
pirate Apr 3, 2026
a32c1bf
more clenaup
pirate Apr 3, 2026
3c4e182
fix custom comments
pirate Apr 3, 2026
a150bbe
add missing keep file
pirate Apr 3, 2026
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
4 changes: 3 additions & 1 deletion examples/full_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,9 @@ async def main() -> None:

# Use the first result
result = results[0]
print(f"Acting on: {result.description}")
if not isinstance(result, dict):
raise RuntimeError(f"Expected observe stream result item to be a dict, got {type(result).__name__}")
print(f"Acting on: {result['description']}")

# Pass the action to Act
act_stream = await session.act(
Expand Down
51 changes: 26 additions & 25 deletions src/stagehand/_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@

import os
from typing import TYPE_CHECKING, Any, Mapping

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Keep the generated client thin: all runtime patch logic lives in `_custom`.
from typing_extensions import Self, Literal, override

import httpx
Expand All @@ -30,9 +33,6 @@
SyncAPIClient,
AsyncAPIClient,
)

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Keep the generated client thin: all runtime patch logic lives in `_custom`.
from ._custom.session import install_stainless_session_patches
from ._custom.sea_server import (
copy_local_mode_kwargs,
Expand All @@ -50,6 +50,7 @@

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
from ._custom.sea_server import SeaServerManager
from .resources.sessions import SessionsResource, AsyncSessionsResource
### </END CUSTOM CODE>

__all__ = [
Expand All @@ -71,11 +72,11 @@


class Stagehand(SyncAPIClient):
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# client options
browserbase_api_key: str | None
browserbase_project_id: str | None
model_api_key: str | None
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# These are assigned indirectly by `configure_client_base_url(...)` so the
# generated class still exposes typed local-mode state for `copy()` and tests.
_server_mode: Literal["remote", "local"]
Expand All @@ -89,6 +90,7 @@ class Stagehand(SyncAPIClient):
_sea_server: SeaServerManager | None
### </END CUSTOM CODE>

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
def __init__(
self,
*,
Expand Down Expand Up @@ -142,7 +144,6 @@ def __init__(

self.model_api_key = model_api_key

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Centralize local-mode state hydration and base-url selection in `_custom`
# so no constructor branching lives in the generated client.
base_url = configure_client_base_url(
Expand All @@ -158,7 +159,7 @@ def __init__(
base_url=base_url,
model_api_key=model_api_key,
)
### </END CUSTOM CODE>
### </END CUSTOM CODE>

super().__init__(
version=__version__,
Expand All @@ -173,29 +174,27 @@ def __init__(

self._default_stream_cls = Stream

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
@override
def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions:
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Start the local SEA server lazily on first request instead of at client
# construction time, then swap the base URL to the started process.
local_base_url = prepare_sync_client_base_url(self)
if local_base_url is not None:
self.base_url = local_base_url
### </END CUSTOM CODE>
return super()._prepare_options(options)

@override
def close(self) -> None:
try:
super().close()
finally:
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Tear down the managed SEA process after HTTP resources close.
close_sync_client_sea_server(self)
### </END CUSTOM CODE>
### </END CUSTOM CODE>

@cached_property
def sessions(self) -> sessions.SessionsResource:
def sessions(self) -> SessionsResource:
from .resources.sessions import SessionsResource

return SessionsResource(self)
Expand All @@ -218,6 +217,7 @@ def qs(self) -> Querystring:
def auth_headers(self) -> dict[str, str]:
return {**self._bb_api_key_auth, **self._bb_project_id_auth, **self._llm_model_api_key_auth}

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
@property
def _bb_api_key_auth(self) -> dict[str, str]:
browserbase_api_key = self.browserbase_api_key
Expand All @@ -243,7 +243,9 @@ def default_headers(self) -> dict[str, str | Omit]:
"X-Stainless-Async": "false",
**self._custom_headers,
}
### </END CUSTOM CODE>

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
def copy(
self,
*,
Expand Down Expand Up @@ -300,7 +302,6 @@ def copy(
max_retries=max_retries if is_given(max_retries) else self.max_retries,
default_headers=headers,
default_query=params,
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Preserve local-mode configuration when cloning the client without
# duplicating that branching logic in generated code.
**copy_local_mode_kwargs(
Expand All @@ -314,12 +315,12 @@ def copy(
local_ready_timeout_s=local_ready_timeout_s,
local_shutdown_on_close=local_shutdown_on_close,
),
### </END CUSTOM CODE>
**_extra_kwargs,
)
### </END CUSTOM CODE>

# Alias for `copy` for nicer inline usage, e.g.
# client.with_options(timeout=10).foo.start(...)
# client.with_options(timeout=10).foo.create(...)
with_options = copy

@override
Expand Down Expand Up @@ -357,11 +358,11 @@ def _make_status_error(


class AsyncStagehand(AsyncAPIClient):
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# client options
browserbase_api_key: str | None
browserbase_project_id: str | None
model_api_key: str | None
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# These are assigned indirectly by `configure_client_base_url(...)` so the
# generated class still exposes typed local-mode state for `copy()` and tests.
_server_mode: Literal["remote", "local"]
Expand All @@ -375,6 +376,7 @@ class AsyncStagehand(AsyncAPIClient):
_sea_server: SeaServerManager | None
### </END CUSTOM CODE>

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
def __init__(
self,
*,
Expand Down Expand Up @@ -428,7 +430,6 @@ def __init__(

self.model_api_key = model_api_key

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Centralize local-mode state hydration and base-url selection in `_custom`
# so no constructor branching lives in the generated client.
base_url = configure_client_base_url(
Expand All @@ -444,7 +445,7 @@ def __init__(
base_url=base_url,
model_api_key=model_api_key,
)
### </END CUSTOM CODE>
### </END CUSTOM CODE>

super().__init__(
version=__version__,
Expand All @@ -459,29 +460,27 @@ def __init__(

self._default_stream_cls = AsyncStream

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
@override
async def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions:
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Start the local SEA server lazily on first request instead of at client
# construction time, then swap the base URL to the started process.
local_base_url = await prepare_async_client_base_url(self)
if local_base_url is not None:
self.base_url = local_base_url
### </END CUSTOM CODE>
return await super()._prepare_options(options)

@override
async def close(self) -> None:
try:
await super().close()
finally:
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Tear down the managed SEA process after HTTP resources close.
await close_async_client_sea_server(self)
### </END CUSTOM CODE>
### </END CUSTOM CODE>

@cached_property
def sessions(self) -> sessions.AsyncSessionsResource:
def sessions(self) -> AsyncSessionsResource:
from .resources.sessions import AsyncSessionsResource

return AsyncSessionsResource(self)
Expand All @@ -504,6 +503,7 @@ def qs(self) -> Querystring:
def auth_headers(self) -> dict[str, str]:
return {**self._bb_api_key_auth, **self._bb_project_id_auth, **self._llm_model_api_key_auth}

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
@property
def _bb_api_key_auth(self) -> dict[str, str]:
browserbase_api_key = self.browserbase_api_key
Expand All @@ -529,7 +529,9 @@ def default_headers(self) -> dict[str, str | Omit]:
"X-Stainless-Async": f"async:{get_async_library()}",
**self._custom_headers,
}
### </END CUSTOM CODE>

### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
def copy(
self,
*,
Expand Down Expand Up @@ -586,7 +588,6 @@ def copy(
max_retries=max_retries if is_given(max_retries) else self.max_retries,
default_headers=headers,
default_query=params,
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
# Preserve local-mode configuration when cloning the client without
# duplicating that branching logic in generated code.
**copy_local_mode_kwargs(
Expand All @@ -600,12 +601,12 @@ def copy(
local_ready_timeout_s=local_ready_timeout_s,
local_shutdown_on_close=local_shutdown_on_close,
),
### </END CUSTOM CODE>
**_extra_kwargs,
)
### </END CUSTOM CODE>

# Alias for `copy` for nicer inline usage, e.g.
# client.with_options(timeout=10).foo.start(...)
# client.with_options(timeout=10).foo.create(...)
with_options = copy

@override
Expand Down
1 change: 1 addition & 0 deletions src/stagehand/_sea/.keep
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@

Loading
Loading