Add AgentDraftSession as private class

src/truefoundry_gateway_sdk/agents/__init__.py

@@ -28,7 +28,7 @@
 from ..types.turn_streaming_event import TurnStreamingEvent
 from ..types.user_tool_approval_event import UserToolApprovalEvent
 from ..types.user_tool_response_event import UserToolResponseEvent
-from .agent_session import AgentSession, AsyncAgentSession
+from .agent_session import AgentSession, AsyncAgentSession, AsyncBaseAgentSession, BaseAgentSession
 from .agent_session_client import AgentSessionClient, AsyncAgentSessionClient
 from .event_delta import DeltaEvents, is_event_delta, merge_event_delta
 from .prepared_turn import AsyncPreparedTurn, PreparedTurn
@@ -42,6 +42,8 @@
     "AgentSessionClient",
     "AsyncAgentSession",
     "AsyncAgentSessionClient",
+    "AsyncBaseAgentSession",
+    "BaseAgentSession",
     "AsyncPreparedTurn",
     "AsyncTurn",
     "DeltaEvents",

src/truefoundry_gateway_sdk/agents/agent_session.py

@@ -21,7 +21,7 @@
 
 def _wrap_turns_pager(
     raw_pager: SyncPager[RawTurn, ListTurnsResponse],
-    session: "AgentSession",
+    session: "BaseAgentSession",
     client: TrueFoundryGateway,
 ) -> SyncPager[Turn, ListTurnsResponse]:
     wrapped_items = [Turn(t, session, client) for t in (raw_pager.items or [])]
@@ -44,7 +44,7 @@ def get_next() -> typing.Optional[SyncPager[Turn, ListTurnsResponse]]:
 
 async def _async_wrap_turns_pager(
     raw_pager: AsyncPager[RawTurn, ListTurnsResponse],
-    session: "AsyncAgentSession",
+    session: "AsyncBaseAgentSession",
     client: AsyncTrueFoundryGateway,
 ) -> AsyncPager[AsyncTurn, ListTurnsResponse]:
     wrapped_items = [AsyncTurn(t, session, client) for t in (raw_pager.items or [])]
@@ -65,23 +65,29 @@ async def get_next() -> typing.Optional[AsyncPager[AsyncTurn, ListTurnsResponse]
     )
 
 
-class AgentSession:
+class BaseAgentSession:
     """
-    A session enriched with convenience methods: prepare_turn, list_turns, get_turn, list_events, cancel.
+    Shared base for AgentSession and AgentDraftSession. Holds common identity fields
+    and provides turn/event convenience methods.
     """
 
-    def __init__(self, session: RawSession, client: TrueFoundryGateway) -> None:
-        self._id: str = session.id
-        self._agent_name: str = session.agent_name
-        self._title: typing.Optional[str] = session.title
-        self._created_by_subject: Subject = session.created_by_subject
-        self._created_at: str = session.created_at
-        self._updated_at: str = session.updated_at
+    def __init__(
+        self,
+        *,
+        id: str,
+        title: typing.Optional[str],
+        created_by_subject: Subject,
+        created_at: str,
+        updated_at: str,
+        client: TrueFoundryGateway,
+    ) -> None:
+        self._id: str = id
+        self._title: typing.Optional[str] = title
+        self._created_by_subject: Subject = created_by_subject
+        self._created_at: str = created_at
+        self._updated_at: str = updated_at
         self._client = client
 
-    def __repr__(self) -> str:
-        return f"AgentSession(id={self._id!r}, agent_name={self._agent_name!r})"
-
     @property
     def id(self) -> str:
         """
@@ -92,16 +98,6 @@ def id(self) -> str:
         """
         return self._id
 
-    @property
-    def agent_name(self) -> str:
-        """
-        Returns
-        -------
-        str
-            Name of the agent for this session.
-        """
-        return self._agent_name
-
     @property
     def title(self) -> typing.Optional[str]:
         """
@@ -274,42 +270,68 @@ def list_events(
         )
 
 
-class AsyncAgentSession:
+class AgentSession(BaseAgentSession):
     """
-    Async version of AgentSession.
+    A session enriched with convenience methods: prepare_turn, list_turns, get_turn, list_events, cancel.
     """
 
-    def __init__(self, session: RawSession, client: AsyncTrueFoundryGateway) -> None:
-        self._id: str = session.id
+    def __init__(self, session: RawSession, client: TrueFoundryGateway) -> None:
+        super().__init__(
+            id=session.id,
+            title=session.title,
+            created_by_subject=session.created_by_subject,
+            created_at=session.created_at,
+            updated_at=session.updated_at,
+            client=client,
+        )
         self._agent_name: str = session.agent_name
-        self._title: typing.Optional[str] = session.title
-        self._created_by_subject: Subject = session.created_by_subject
-        self._created_at: str = session.created_at
-        self._updated_at: str = session.updated_at
-        self._client = client
 
     def __repr__(self) -> str:
-        return f"AsyncAgentSession(id={self._id!r}, agent_name={self._agent_name!r})"
+        return f"AgentSession(id={self._id!r}, agent_name={self._agent_name!r})"
 
     @property
-    def id(self) -> str:
+    def agent_name(self) -> str:
         """
         Returns
         -------
         str
-            Unique identifier of this session.
+            Name of the agent for this session.
         """
-        return self._id
+        return self._agent_name
+
+
+class AsyncBaseAgentSession:
+    """
+    Async shared base for AsyncAgentSession and AsyncAgentDraftSession. Holds common identity
+    fields and provides turn/event convenience methods.
+    """
+
+    def __init__(
+        self,
+        *,
+        id: str,
+        title: typing.Optional[str],
+        created_by_subject: Subject,
+        created_at: str,
+        updated_at: str,
+        client: AsyncTrueFoundryGateway,
+    ) -> None:
+        self._id: str = id
+        self._title: typing.Optional[str] = title
+        self._created_by_subject: Subject = created_by_subject
+        self._created_at: str = created_at
+        self._updated_at: str = updated_at
+        self._client = client
 
     @property
-    def agent_name(self) -> str:
+    def id(self) -> str:
         """
         Returns
         -------
         str
-            Name of the agent for this session.
+            Unique identifier of this session.
         """
-        return self._agent_name
+        return self._id
 
     @property
     def title(self) -> typing.Optional[str]:
@@ -481,3 +503,33 @@ async def list_events(
             limit=limit,
             request_options=request_options,
         )
+
+
+class AsyncAgentSession(AsyncBaseAgentSession):
+    """
+    Async version of AgentSession.
+    """
+
+    def __init__(self, session: RawSession, client: AsyncTrueFoundryGateway) -> None:
+        super().__init__(
+            id=session.id,
+            title=session.title,
+            created_by_subject=session.created_by_subject,
+            created_at=session.created_at,
+            updated_at=session.updated_at,
+            client=client,
+        )
+        self._agent_name: str = session.agent_name
+
+    def __repr__(self) -> str:
+        return f"AsyncAgentSession(id={self._id!r}, agent_name={self._agent_name!r})"
+
+    @property
+    def agent_name(self) -> str:
+        """
+        Returns
+        -------
+        str
+            Name of the agent for this session.
+        """
+        return self._agent_name

src/truefoundry_gateway_sdk/agents/prepared_turn.py

@@ -19,7 +19,7 @@
     from ..types.turn_event import TurnEvent
     from ..types.turn_input_item import TurnInputItem
     from ..types.turn_state import TurnState
-    from .agent_session import AgentSession, AsyncAgentSession
+    from .agent_session import AsyncBaseAgentSession, BaseAgentSession
 
 
 class PreparedTurn:
@@ -33,12 +33,12 @@ def __init__(
         *,
         input: typing.Optional[typing.Sequence[TurnInputItem]],
         previous_turn_id: typing.Optional[PreviousTurnIdInput],
-        session: "AgentSession",
+        session: "BaseAgentSession",
         client: TrueFoundryGateway,
     ) -> None:
         self._input = list(input) if input is not None else None
         self._previous_turn_id = previous_turn_id
-        self._session: "AgentSession" = session
+        self._session: "BaseAgentSession" = session
         self._session_id: str = session.id
         self._client = client
         self._started: bool = False
@@ -50,11 +50,11 @@ def __repr__(self) -> str:
     # --- Properties that delegate to the inner Turn (None until execute is called) ---
 
     @property
-    def session(self) -> "AgentSession":
+    def session(self) -> "BaseAgentSession":
         """
         Returns
         -------
-        AgentSession
+        BaseAgentSession
             Parent session this prepared turn belongs to.
         """
         return self._session
@@ -385,12 +385,12 @@ def __init__(
         *,
         input: typing.Optional[typing.Sequence[TurnInputItem]],
         previous_turn_id: typing.Optional[PreviousTurnIdInput],
-        session: "AsyncAgentSession",
+        session: "AsyncBaseAgentSession",
         client: AsyncTrueFoundryGateway,
     ) -> None:
         self._input = list(input) if input is not None else None
         self._previous_turn_id = previous_turn_id
-        self._session: "AsyncAgentSession" = session
+        self._session: "AsyncBaseAgentSession" = session
         self._session_id: str = session.id
         self._client = client
         self._started: bool = False
@@ -402,11 +402,11 @@ def __repr__(self) -> str:
     # --- Properties that delegate to the inner AsyncTurn (None until execute is called) ---
 
     @property
-    def session(self) -> "AsyncAgentSession":
+    def session(self) -> "AsyncBaseAgentSession":
         """
         Returns
         -------
-        AsyncAgentSession
+        AsyncBaseAgentSession
             Parent session this prepared turn belongs to.
         """
         return self._session

src/truefoundry_gateway_sdk/agents/private/__init__.py

@@ -0,0 +1,3 @@
+from .agent_draft_session import AgentDraftSession, AsyncAgentDraftSession
+
+__all__ = ["AgentDraftSession", "AsyncAgentDraftSession"]