centralize custom code

Author: pirate

CONTRIBUTING.md

@@ -35,7 +35,7 @@ The SDK supports running a local Stagehand server for development and testing. T
 Run the download script to automatically download the correct binary:
 
 ```sh
-$ uv run python scripts/download-binary.py
+$ uv run python scripts/download_binary.py
 ```
 
 This will:
@@ -64,7 +64,7 @@ Instead of placing the binary in `bin/sea/`, you can point to any binary locatio
 
 ```sh
 $ export STAGEHAND_SEA_BINARY=/path/to/your/stagehand-binary
-$ uv run python test_local_mode.py
+$ uv run python scripts/test_local_mode.py
 ```
 
 ## Adding and running examples

pyproject.toml

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

scripts/download-binary.py scripts/download_binary.pyscripts/download-binary.py renamed to scripts/download_binary.py

@@ -6,11 +6,11 @@
 and places it in bin/sea/ for use during development and testing.
 
 Usage:
-    python scripts/download-binary.py [--version VERSION]
+    python scripts/download_binary.py [--version VERSION]
 
 Examples:
-    python scripts/download-binary.py
-    python scripts/download-binary.py --version v3.2.0
+    python scripts/download_binary.py
+    python scripts/download_binary.py --version v3.2.0
 """
 from __future__ import annotations
 
@@ -179,7 +179,7 @@ def reporthook(block_num: int, block_size: int, total_size: int) -> None:
 
         size_mb = dest_path.stat().st_size / (1024 * 1024)
         print(f"✅ Downloaded successfully: {dest_path} ({size_mb:.1f} MB)")
-        print(f"\n💡 You can now run: uv run python test_local_mode.py")
+        print("\n💡 You can now run: uv run python scripts/test_local_mode.py")
 
     except urllib.error.HTTPError as e:  # type: ignore[misc]
         print(f"\n❌ Error: Failed to download binary (HTTP {e.code})")  # type: ignore[union-attr]
@@ -197,9 +197,9 @@ def main() -> None:
         formatter_class=argparse.RawDescriptionHelpFormatter,
         epilog="""
 Examples:
-  python scripts/download-binary.py
-  python scripts/download-binary.py --version v3.2.0
-  python scripts/download-binary.py --version stagehand-server-v3/v3.2.0
+  python scripts/download_binary.py
+  python scripts/download_binary.py --version v3.2.0
+  python scripts/download_binary.py --version stagehand-server-v3/v3.2.0
         """,
     )
     parser.add_argument(

scripts/test_local_mode.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""Quick test of local server mode with the embedded binary."""
+
+import os
+import sys
+import traceback
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
+
+from stagehand import Stagehand
+
+
+def main() -> None:
+    model_api_key = os.environ.get("MODEL_API_KEY") or os.environ.get("OPENAI_API_KEY")
+    if not model_api_key:
+        print("❌ Error: MODEL_API_KEY or OPENAI_API_KEY environment variable not set")  # noqa: T201
+        print("   Set it with: export MODEL_API_KEY='sk-proj-...'")  # noqa: T201
+        sys.exit(1)
+
+    os.environ["BROWSERBASE_FLOW_LOGS"] = "1"
+
+    print("🚀 Testing local server mode...")  # noqa: T201
+    client = None
+
+    try:
+        print("📦 Creating Stagehand client in local mode...")  # noqa: T201
+        client = Stagehand(
+            server="local",
+            browserbase_api_key="local",
+            browserbase_project_id="local",
+            model_api_key=model_api_key,
+            local_headless=True,
+            local_port=0,
+            local_ready_timeout_s=15.0,
+        )
+
+        print("🔧 Starting session (this will start the local server)...")  # noqa: T201
+        session = client.sessions.start(
+            model_name="openai/gpt-5-nano",
+            browser={  # type: ignore[arg-type]
+                "type": "local",
+                "launchOptions": {},
+            },
+        )
+        session_id = session.data.session_id
+
+        print(f"✅ Session started: {session_id}")  # noqa: T201
+        print(f"🌐 Server running at: {client.base_url}")  # noqa: T201
+
+        print("\n📍 Navigating to example.com...")  # noqa: T201
+        client.sessions.navigate(id=session_id, url="https://example.com")
+        print("✅ Navigation complete")  # noqa: T201
+
+        print("\n🔍 Extracting page heading...")  # noqa: T201
+        result = client.sessions.extract(
+            id=session_id,
+            instruction="Extract the main heading text from the page",
+        )
+        print(f"📄 Extracted: {result.data.result}")  # noqa: T201
+
+        print("\n🛑 Ending session...")  # noqa: T201
+        client.sessions.end(id=session_id)
+        print("✅ Session ended")  # noqa: T201
+        print("\n🎉 All tests passed!")  # noqa: T201
+    except Exception as exc:
+        print(f"\n❌ Error: {exc}")  # noqa: T201
+        traceback.print_exc()
+        sys.exit(1)
+    finally:
+        if client is not None:
+            print("\n🔌 Closing client (will shut down server)...")  # noqa: T201
+            client.close()
+            print("✅ Server shut down successfully!")  # noqa: T201
+
+
+if __name__ == "__main__":
+    main()

src/stagehand/__init__.py

@@ -38,6 +38,7 @@
 )
 from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
 from ._utils._logs import setup_logging as _setup_logging
+from ._custom.session import Session, AsyncSession
 
 __all__ = [
     "types",
@@ -73,6 +74,8 @@
     "AsyncStream",
     "Stagehand",
     "AsyncStagehand",
+    "Session",
+    "AsyncSession",
     "file_from_path",
     "BaseModel",
     "DEFAULT_TIMEOUT",