Skip to content

Commit ff29af4

Browse files
committed
Centralize custom code
1 parent fc92e13 commit ff29af4

File tree

15 files changed

+944
-633
lines changed

15 files changed

+944
-633
lines changed

CONTRIBUTING.md

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ uv run python script.py
2424

2525
Most of the SDK is generated code. Modifications to code will be persisted between generations, but may
2626
result in merge conflicts between manual patches and changes from the generator. The generator will never
27-
modify the contents of the `src/stagehand/lib/` and `examples/` directories.
27+
modify the contents of the `src/stagehand/_custom/` and `examples/` directories.
2828

2929
## Setting up the local server binary (for development)
3030

@@ -35,7 +35,7 @@ The SDK supports running a local Stagehand server for development and testing. T
3535
Run the download script to automatically download the correct binary:
3636

3737
```sh
38-
$ uv run python scripts/download-binary.py
38+
$ uv run python scripts/download_binary.py
3939
```
4040

4141
This will:
@@ -64,7 +64,7 @@ Instead of placing the binary in `bin/sea/`, you can point to any binary locatio
6464

6565
```sh
6666
$ export STAGEHAND_SEA_BINARY=/path/to/your/stagehand-binary
67-
$ uv run python test_local_mode.py
67+
$ uv run python scripts/test_local_mode.py
6868
```
6969

7070
## Adding and running examples

pyproject.toml

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,6 @@ exclude = [
137137
"hatch_build.py",
138138
"examples",
139139
"scripts",
140-
"test_local_mode.py",
141140
]
142141

143142
reportImplicitOverride = true
@@ -156,7 +155,7 @@ show_error_codes = true
156155
#
157156
# We also exclude our `tests` as mypy doesn't always infer
158157
# types correctly and Pyright will still catch any type errors.
159-
exclude = ['src/stagehand/_files.py', '_dev/.*.py', 'tests/.*', 'hatch_build.py', 'examples/.*', 'scripts/.*', 'test_local_mode.py']
158+
exclude = ['src/stagehand/_files.py', '_dev/.*.py', 'tests/.*', 'hatch_build.py', 'examples/.*', 'scripts/.*']
160159

161160
strict_equality = true
162161
implicit_reexport = true

scripts/download_binary.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@
66
and places it in bin/sea/ for use during development and testing.
77
88
Usage:
9-
python scripts/download-binary.py [--version VERSION]
9+
python scripts/download_binary.py [--version VERSION]
1010
1111
Examples:
12-
python scripts/download-binary.py
13-
python scripts/download-binary.py --version v3.2.0
12+
python scripts/download_binary.py
13+
python scripts/download_binary.py --version v3.2.0
1414
"""
1515
from __future__ import annotations
1616

@@ -179,7 +179,7 @@ def reporthook(block_num: int, block_size: int, total_size: int) -> None:
179179

180180
size_mb = dest_path.stat().st_size / (1024 * 1024)
181181
print(f"✅ Downloaded successfully: {dest_path} ({size_mb:.1f} MB)")
182-
print(f"\n💡 You can now run: uv run python test_local_mode.py")
182+
print("\n💡 You can now run: uv run python scripts/test_local_mode.py")
183183

184184
except urllib.error.HTTPError as e: # type: ignore[misc]
185185
print(f"\n❌ Error: Failed to download binary (HTTP {e.code})") # type: ignore[union-attr]
@@ -197,9 +197,9 @@ def main() -> None:
197197
formatter_class=argparse.RawDescriptionHelpFormatter,
198198
epilog="""
199199
Examples:
200-
python scripts/download-binary.py
201-
python scripts/download-binary.py --version v3.2.0
202-
python scripts/download-binary.py --version stagehand-server-v3/v3.2.0
200+
python scripts/download_binary.py
201+
python scripts/download_binary.py --version v3.2.0
202+
python scripts/download_binary.py --version stagehand-server-v3/v3.2.0
203203
""",
204204
)
205205
parser.add_argument(

scripts/test_local_mode.py

Lines changed: 59 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -3,72 +3,76 @@
33

44
import os
55
import sys
6+
import traceback
7+
from pathlib import Path
68

7-
# Add src to path for local testing
8-
sys.path.insert(0, "src")
9+
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
910

1011
from stagehand import Stagehand
1112

12-
# Set required API key for LLM operations
13-
if not os.environ.get("MODEL_API_KEY") and not os.environ.get("OPENAI_API_KEY"):
14-
print("❌ Error: MODEL_API_KEY or OPENAI_API_KEY environment variable not set") # noqa: T201
15-
print(" Set it with: export MODEL_API_KEY='sk-proj-...'") # noqa: T201
16-
sys.exit(1)
1713

18-
print("🚀 Testing local server mode...") # noqa: T201
14+
def main() -> None:
15+
model_api_key = os.environ.get("MODEL_API_KEY") or os.environ.get("OPENAI_API_KEY")
16+
if not model_api_key:
17+
print("❌ Error: MODEL_API_KEY or OPENAI_API_KEY environment variable not set") # noqa: T201
18+
print(" Set it with: export MODEL_API_KEY='sk-proj-...'") # noqa: T201
19+
sys.exit(1)
1920

20-
try:
21-
# Create client in local mode - will use bundled binary
22-
print("📦 Creating Stagehand client in local mode...") # noqa: T201
23-
client = Stagehand(
24-
server="local",
25-
browserbase_api_key="local", # Dummy value - not used in local mode
26-
browserbase_project_id="local", # Dummy value - not used in local mode
27-
model_api_key=os.environ.get("MODEL_API_KEY") or os.environ["OPENAI_API_KEY"],
28-
local_headless=True,
29-
local_port=0, # Auto-pick free port
30-
local_ready_timeout_s=15.0, # Give it time to start
31-
)
21+
os.environ["BROWSERBASE_FLOW_LOGS"] = "1"
3222

33-
print("🔧 Starting session (this will start the local server)...") # noqa: T201
34-
session = client.sessions.start(
35-
model_name="openai/gpt-5-nano",
36-
browser={ # type: ignore[arg-type]
37-
"type": "local",
38-
"launchOptions": {}, # Launch local Playwright browser with defaults
39-
},
40-
)
41-
session_id = session.data.session_id
23+
print("🚀 Testing local server mode...") # noqa: T201
24+
client = None
4225

43-
print(f"✅ Session started: {session_id}") # noqa: T201
44-
print(f"🌐 Server running at: {client.base_url}") # noqa: T201
26+
try:
27+
print("📦 Creating Stagehand client in local mode...") # noqa: T201
28+
client = Stagehand(
29+
server="local",
30+
browserbase_api_key="local",
31+
browserbase_project_id="local",
32+
model_api_key=model_api_key,
33+
local_headless=True,
34+
local_port=0,
35+
local_ready_timeout_s=15.0,
36+
)
4537

46-
print("\n📍 Navigating to example.com...") # noqa: T201
47-
client.sessions.navigate(
48-
id=session_id,
49-
url="https://example.com",
50-
)
51-
print("✅ Navigation complete") # noqa: T201
38+
print("🔧 Starting session (this will start the local server)...") # noqa: T201
39+
session = client.sessions.start(
40+
model_name="openai/gpt-5-nano",
41+
browser={ # type: ignore[arg-type]
42+
"type": "local",
43+
"launchOptions": {},
44+
},
45+
)
46+
session_id = session.data.session_id
5247

53-
print("\n🔍 Extracting page heading...") # noqa: T201
54-
result = client.sessions.extract(
55-
id=session_id,
56-
instruction="Extract the main heading text from the page",
57-
)
58-
print(f"📄 Extracted: {result.data.result}") # noqa: T201
48+
print(f"✅ Session started: {session_id}") # noqa: T201
49+
print(f"🌐 Server running at: {client.base_url}") # noqa: T201
5950

60-
print("\n🛑 Ending session...") # noqa: T201
61-
client.sessions.end(id=session_id)
62-
print("✅ Session ended") # noqa: T201
51+
print("\n📍 Navigating to example.com...") # noqa: T201
52+
client.sessions.navigate(id=session_id, url="https://example.com")
53+
print("✅ Navigation complete") # noqa: T201
6354

64-
print("\n🔌 Closing client (will shut down server)...") # noqa: T201
65-
client.close()
66-
print("✅ Server shut down successfully!") # noqa: T201
55+
print("\n🔍 Extracting page heading...") # noqa: T201
56+
result = client.sessions.extract(
57+
id=session_id,
58+
instruction="Extract the main heading text from the page",
59+
)
60+
print(f"📄 Extracted: {result.data.result}") # noqa: T201
6761

68-
print("\n🎉 All tests passed!") # noqa: T201
62+
print("\n🛑 Ending session...") # noqa: T201
63+
client.sessions.end(id=session_id)
64+
print("✅ Session ended") # noqa: T201
65+
print("\n🎉 All tests passed!") # noqa: T201
66+
except Exception as exc:
67+
print(f"\n❌ Error: {exc}") # noqa: T201
68+
traceback.print_exc()
69+
sys.exit(1)
70+
finally:
71+
if client is not None:
72+
print("\n🔌 Closing client (will shut down server)...") # noqa: T201
73+
client.close()
74+
print("✅ Server shut down successfully!") # noqa: T201
6975

70-
except Exception as e:
71-
print(f"\n❌ Error: {e}") # noqa: T201
72-
import traceback
73-
traceback.print_exc()
74-
sys.exit(1)
76+
77+
if __name__ == "__main__":
78+
main()

src/stagehand/__init__.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,13 @@
3939
from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
4040
from ._utils._logs import setup_logging as _setup_logging
4141

42+
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
43+
# Re-export the public bound session types from `_custom` so users can type
44+
# against `stagehand.Session` instead of importing from private modules.
45+
from ._custom.session import Session, AsyncSession
46+
47+
### </END CUSTOM CODE>
48+
4249
__all__ = [
4350
"types",
4451
"__version__",
@@ -73,6 +80,10 @@
7380
"AsyncStream",
7481
"Stagehand",
7582
"AsyncStagehand",
83+
### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
84+
"Session",
85+
"AsyncSession",
86+
### </END CUSTOM CODE>
7687
"file_from_path",
7788
"BaseModel",
7889
"DEFAULT_TIMEOUT",

0 commit comments

Comments
 (0)