AGE-1320: Add internal list sessions to SDK

src/truefoundry_gateway_sdk/agents/__init__.py

@@ -32,17 +32,28 @@
 from .agent_session_client import AgentSessionClient, AsyncAgentSessionClient
 from .event_delta import DeltaEvents, is_event_delta, merge_event_delta
 from .prepared_turn import AsyncPreparedTurn, PreparedTurn
+from .private import (
+    AgentDraftSession,
+    AsyncAgentDraftSession,
+    AsyncPrivateAgentSessionClient,
+    PrivateAgentSessionClient,
+)
+from .session_mixin import AsyncSessionMixin, SessionMixin
 from .turn import AsyncTurn, Turn
 from .turn_stream_data import TurnStreamData
 
 __all__ = [  # type: ignore[reportUnsupportedDunderAll]
     *_private_agents.__all__,
     "ActionRequiredEvent",
+    "AgentDraftSession",
     "AgentSession",
     "AgentSessionClient",
+    "AsyncAgentDraftSession",
     "AsyncAgentSession",
     "AsyncAgentSessionClient",
     "AsyncPreparedTurn",
+    "AsyncPrivateAgentSessionClient",
+    "AsyncSessionMixin",
     "AsyncTurn",
     "DeltaEvents",
     "McpAuthRequiredEvent",
@@ -52,7 +63,9 @@
     "ModelMessageUsage",
     "ModelMessageUsageInputTokensBreakdown",
     "PreparedTurn",
+    "PrivateAgentSessionClient",
     "SandboxCreatedEvent",
+    "SessionMixin",
     "ThreadCreatedEvent",
     "ThreadDoneEvent",
     "ToolApprovalRequiredEvent",

src/truefoundry_gateway_sdk/agents/agent_session.py

@@ -2,72 +2,27 @@
 
 import typing
 
-from ..core.pagination import AsyncPager, SyncPager
-from .prepared_turn import AsyncPreparedTurn, PreparedTurn
-from .turn import AsyncTurn, Turn
+from .session_mixin import AsyncSessionMixin, SessionMixin
 
 if typing.TYPE_CHECKING:
     from ..client import AsyncTrueFoundryGateway, TrueFoundryGateway
+    from ..core.pagination import AsyncPager, SyncPager
     from ..core.request_options import RequestOptions
     from ..types.list_session_events_response import ListSessionEventsResponse
     from ..types.list_turns_response import ListTurnsResponse
     from ..types.previous_turn_id_input import PreviousTurnIdInput
     from ..types.session import Session as RawSession
     from ..types.session_event_item import SessionEventItem
     from ..types.subject import Subject
-    from ..types.turn import Turn as RawTurn
     from ..types.turn_input_item import TurnInputItem
-
-
-def _wrap_turns_pager(
-    raw_pager: SyncPager[RawTurn, ListTurnsResponse],
-    session: "AgentSession",
-    client: TrueFoundryGateway,
-) -> SyncPager[Turn, ListTurnsResponse]:
-    wrapped_items = [Turn(t, session, client) for t in (raw_pager.items or [])]
-
-    def get_next() -> typing.Optional[SyncPager[Turn, ListTurnsResponse]]:
-        if raw_pager.get_next is None:
-            return None
-        next_raw = raw_pager.get_next()
-        if next_raw is None:
-            return None
-        return _wrap_turns_pager(next_raw, session, client)
-
-    return SyncPager(
-        get_next=get_next if raw_pager.has_next else None,
-        has_next=raw_pager.has_next,
-        items=wrapped_items,
-        response=raw_pager.response,
-    )
-
-
-async def _async_wrap_turns_pager(
-    raw_pager: AsyncPager[RawTurn, ListTurnsResponse],
-    session: "AsyncAgentSession",
-    client: AsyncTrueFoundryGateway,
-) -> AsyncPager[AsyncTurn, ListTurnsResponse]:
-    wrapped_items = [AsyncTurn(t, session, client) for t in (raw_pager.items or [])]
-
-    async def get_next() -> typing.Optional[AsyncPager[AsyncTurn, ListTurnsResponse]]:
-        if raw_pager.get_next is None:
-            return None
-        next_raw = await raw_pager.get_next()
-        if next_raw is None:
-            return None
-        return await _async_wrap_turns_pager(next_raw, session, client)
-
-    return AsyncPager(
-        get_next=get_next if raw_pager.has_next else None,
-        has_next=raw_pager.has_next,
-        items=wrapped_items,
-        response=raw_pager.response,
-    )
+    from .prepared_turn import AsyncPreparedTurn, PreparedTurn
+    from .turn import AsyncTurn, Turn
 
 
 class AgentSession:
     """
     A session enriched with convenience methods: prepare_turn, list_turns, get_turn, list_events, cancel.
+    Turn operations are delegated to a shared :class:`SessionMixin`.
     """
 
     def __init__(self, session: RawSession, client: TrueFoundryGateway) -> None:
@@ -77,11 +32,21 @@ def __init__(self, session: RawSession, client: TrueFoundryGateway) -> None:
         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
+        self._mixin = SessionMixin(session.id, client)
 
     def __repr__(self) -> str:
         return f"AgentSession(id={self._id!r}, agent_name={self._agent_name!r})"
 
+    @property
+    def type(self) -> typing.Literal["session"]:
+        """
+        Returns
+        -------
+        typing.Literal["session"]
+            Discriminant distinguishing a saved session from a draft session.
+        """
+        return "session"
+
     @property
     def id(self) -> str:
         """
@@ -163,12 +128,7 @@ def prepare_turn(
         PreparedTurn
             Staged turn.
         """
-        return PreparedTurn(
-            input=input,
-            previous_turn_id=previous_turn_id,
-            session=self,
-            client=self._client,
-        )
+        return self._mixin.prepare_turn(self, input=input, previous_turn_id=previous_turn_id)
 
     def list_turns(
         self,
@@ -194,10 +154,7 @@ def list_turns(
         SyncPager[Turn, ListTurnsResponse]
             Paginated turns.
         """
-        raw_pager = self._client.agents.sessions.list_turns(
-            self._id, page_token=page_token, limit=limit, request_options=request_options
-        )
-        return _wrap_turns_pager(raw_pager, self, self._client)
+        return self._mixin.list_turns(self, page_token=page_token, limit=limit, request_options=request_options)
 
     def get_turn(
         self,
@@ -220,8 +177,7 @@ def get_turn(
         Turn
             Turn data.
         """
-        response = self._client.agents.sessions.get_turn(self._id, turn_id, request_options=request_options)
-        return Turn(response.data, self, self._client)
+        return self._mixin.get_turn(self, turn_id, request_options=request_options)
 
     def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None:
         """
@@ -236,7 +192,7 @@ def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) ->
         -------
         None
         """
-        self._client.agents.sessions.cancel(self._id, request_options=request_options)
+        self._mixin.cancel(request_options=request_options)
 
     def list_events(
         self,
@@ -265,8 +221,7 @@ def list_events(
         SyncPager[SessionEventItem, ListSessionEventsResponse]
             Paginated session events.
         """
-        return self._client.agents.sessions.list_events(
-            self._id,
+        return self._mixin.list_events(
             page_token=page_token,
             last_turn_id=last_turn_id,
             limit=limit,
@@ -286,11 +241,21 @@ def __init__(self, session: RawSession, client: AsyncTrueFoundryGateway) -> None
         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
+        self._mixin = AsyncSessionMixin(session.id, client)
 
     def __repr__(self) -> str:
         return f"AsyncAgentSession(id={self._id!r}, agent_name={self._agent_name!r})"
 
+    @property
+    def type(self) -> typing.Literal["session"]:
+        """
+        Returns
+        -------
+        typing.Literal["session"]
+            Discriminant distinguishing a saved session from a draft session.
+        """
+        return "session"
+
     @property
     def id(self) -> str:
         """
@@ -372,12 +337,7 @@ def prepare_turn(
         AsyncPreparedTurn
             Staged turn.
         """
-        return AsyncPreparedTurn(
-            input=input,
-            previous_turn_id=previous_turn_id,
-            session=self,
-            client=self._client,
-        )
+        return self._mixin.prepare_turn(self, input=input, previous_turn_id=previous_turn_id)
 
     async def list_turns(
         self,
@@ -403,10 +363,7 @@ async def list_turns(
         AsyncPager[AsyncTurn, ListTurnsResponse]
             Paginated turns.
         """
-        raw_pager = await self._client.agents.sessions.list_turns(
-            self._id, page_token=page_token, limit=limit, request_options=request_options
-        )
-        return await _async_wrap_turns_pager(raw_pager, self, self._client)
+        return await self._mixin.list_turns(self, page_token=page_token, limit=limit, request_options=request_options)
 
     async def get_turn(
         self,
@@ -429,8 +386,7 @@ async def get_turn(
         AsyncTurn
             Turn data.
         """
-        response = await self._client.agents.sessions.get_turn(self._id, turn_id, request_options=request_options)
-        return AsyncTurn(response.data, self, self._client)
+        return await self._mixin.get_turn(self, turn_id, request_options=request_options)
 
     async def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None:
         """
@@ -445,7 +401,7 @@ async def cancel(self, *, request_options: typing.Optional[RequestOptions] = Non
         -------
         None
         """
-        await self._client.agents.sessions.cancel(self._id, request_options=request_options)
+        await self._mixin.cancel(request_options=request_options)
 
     async def list_events(
         self,
@@ -474,8 +430,7 @@ async def list_events(
         AsyncPager[SessionEventItem, ListSessionEventsResponse]
             Paginated session events.
         """
-        return await self._client.agents.sessions.list_events(
-            self._id,
+        return await self._mixin.list_events(
             page_token=page_token,
             last_turn_id=last_turn_id,
             limit=limit,

src/truefoundry_gateway_sdk/agents/prepared_turn.py

@@ -20,6 +20,7 @@
     from ..types.turn_input_item import TurnInputItem
     from ..types.turn_state import TurnState
     from .agent_session import AgentSession, AsyncAgentSession
+    from .private.agent_draft_session import AgentDraftSession, AsyncAgentDraftSession
 
 
 class PreparedTurn:
@@ -33,12 +34,12 @@ def __init__(
         *,
         input: typing.Optional[typing.Sequence[TurnInputItem]],
         previous_turn_id: typing.Optional[PreviousTurnIdInput],
-        session: "AgentSession",
+        session: typing.Union["AgentSession", "AgentDraftSession"],
         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: typing.Union["AgentSession", "AgentDraftSession"] = session
         self._session_id: str = session.id
         self._client = client
         self._started: bool = False
@@ -50,11 +51,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) -> typing.Union["AgentSession", "AgentDraftSession"]:
         """
         Returns
         -------
-        AgentSession
+        typing.Union[AgentSession, AgentDraftSession]
             Parent session this prepared turn belongs to.
         """
         return self._session
@@ -385,12 +386,12 @@ def __init__(
         *,
         input: typing.Optional[typing.Sequence[TurnInputItem]],
         previous_turn_id: typing.Optional[PreviousTurnIdInput],
-        session: "AsyncAgentSession",
+        session: typing.Union["AsyncAgentSession", "AsyncAgentDraftSession"],
         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: typing.Union["AsyncAgentSession", "AsyncAgentDraftSession"] = session
         self._session_id: str = session.id
         self._client = client
         self._started: bool = False
@@ -402,11 +403,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) -> typing.Union["AsyncAgentSession", "AsyncAgentDraftSession"]:
         """
         Returns
         -------
-        AsyncAgentSession
+        typing.Union[AsyncAgentSession, AsyncAgentDraftSession]
             Parent session this prepared turn belongs to.
         """
         return self._session

src/truefoundry_gateway_sdk/agents/private/__init__.py

@@ -0,0 +1,9 @@
+from .agent_draft_session import AgentDraftSession, AsyncAgentDraftSession
+from .private_agent_session_client import AsyncPrivateAgentSessionClient, PrivateAgentSessionClient
+
+__all__ = [
+    "AgentDraftSession",
+    "AsyncAgentDraftSession",
+    "AsyncPrivateAgentSessionClient",
+    "PrivateAgentSessionClient",
+]