make AsyncSession and Session subclass SessionStartResponse for compat with autogen sdk docs

pirate · pirate · commit 3bcc3a919378 · 2026-01-12T15:13:54.000-08:00
diff --git a/src/stagehand/resources/sessions_helpers.py b/src/stagehand/resources/sessions_helpers.py
@@ -59,7 +59,7 @@ def create(
             extra_body=extra_body,
             timeout=timeout,
         )
-        return Session(self._client, start_response.data.session_id)
+        return Session(self._client, start_response.data.session_id, data=start_response.data, success=start_response.success)
 
 
 class AsyncSessionsResourceWithHelpers(AsyncSessionsResource):
@@ -86,7 +86,7 @@ async def create(
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> AsyncSession:
-        start_response = await self.start(
+        start_response: SessionStartResponse = await self.start(
             model_name=model_name,
             act_timeout_ms=act_timeout_ms,
             browser=browser,
@@ -107,4 +107,4 @@ async def create(
             extra_body=extra_body,
             timeout=timeout,
         )
-        return AsyncSession(self._client, start_response.data.session_id)
+        return AsyncSession(self._client, start_response.data.session_id, data=start_response.data, success=start_response.success)
diff --git a/src/stagehand/session.py b/src/stagehand/session.py
@@ -16,6 +16,7 @@
     session_navigate_params,
 )
 from ._types import Body, Omit, Query, Headers, NotGiven, omit, not_given
+from .types.session_start_response import SessionStartResponse, Data as SessionStartResponseData
 from .types.session_act_response import SessionActResponse
 from .types.session_end_response import SessionEndResponse
 from .types.session_execute_response import SessionExecuteResponse
@@ -36,12 +37,15 @@ def _with_default_frame_id(params: TSessionParams) -> TSessionParams:
     from ._client import Stagehand, AsyncStagehand
 
 
-class Session:
+class Session(SessionStartResponse):
     """A Stagehand session bound to a specific `session_id`."""
 
-    def __init__(self, client: Stagehand, id: str) -> None:
+    def __init__(self, client: Stagehand, id: str, data: SessionStartResponseData, success: bool) -> None:
         self._client = client
         self.id = id
+        # in case user tries to use client.sessions.start(...) return value as a SessionStartResponse dataclass/dict
+        super().__init__(data=data, success=success)
+    
 
     def navigate(
         self,
@@ -158,12 +162,13 @@ def end(
         )
 
 
-class AsyncSession:
+class AsyncSession(SessionStartResponse):
     """Async variant of `Session`."""
 
-    def __init__(self, client: AsyncStagehand, id: str) -> None:
+    def __init__(self, client: AsyncStagehand, id: str, data: SessionStartResponseData, success: bool) -> None:
         self._client = client
         self.id = id
+        super().__init__(data=data, success=success)
 
     async def navigate(
         self,
