Centralize custom code (#332)

* Move custom code and scripts

* Centralize custom code

* add compat

* always detatch cdp

Author: pirate

CONTRIBUTING.md

@@ -24,7 +24,7 @@ uv run python script.py
 
 Most of the SDK is generated code. Modifications to code will be persisted between generations, but may
 result in merge conflicts between manual patches and changes from the generator. The generator will never
-modify the contents of the `src/stagehand/lib/` and `examples/` directories.
+modify the contents of the `src/stagehand/_custom/` and `examples/` directories.
 
 ## Setting up the local server binary (for development)
 
@@ -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

@@ -39,6 +39,13 @@
 from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
 from ._utils._logs import setup_logging as _setup_logging
 
+### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
+# Re-export the public bound session types from `_custom` so users can type
+# against `stagehand.Session` instead of importing from private modules.
+from ._custom.session import Session, AsyncSession
+
+### </END CUSTOM CODE>
+
 __all__ = [
     "types",
     "__version__",
@@ -73,6 +80,10 @@
     "AsyncStream",
     "Stagehand",
     "AsyncStagehand",
+    ### <CUSTOM CODE HANDWRITTEN BY STAGEHAND TEAM (not codegen)>
+    "Session",
+    "AsyncSession",
+    ### </END CUSTOM CODE>
     "file_from_path",
     "BaseModel",
     "DEFAULT_TIMEOUT",