Skip to content

Add search sessions api to private client #12

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chiragjn merged 1 commit into from
Jul 24, 2026
Merged

Add search sessions api to private client #12

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
249 changes: 192 additions & 57 deletions src/truefoundry_gateway_sdk/agents/private/private_agent_session_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,63 +14,48 @@
from ...core.logging import LogConfig, Logger
from ...core.request_options import RequestOptions
from ...types.agent_spec import AgentSpec
from ...types.created_by_subject_type import CreatedBySubjectType
from ...types.draft_session import DraftSession as RawDraftSession
from ...types.list_draft_sessions_order import ListDraftSessionsOrder
from ...types.list_draft_sessions_response import ListDraftSessionsResponse
from ...types.list_owned_sessions_order import ListOwnedSessionsOrder
from ...types.list_owned_sessions_response import ListOwnedSessionsResponse
from ...types.list_owned_sessions_response_data_item import ListOwnedSessionsResponseDataItem
from ...types.search_sessions_order import SearchSessionsOrder
from ...types.search_sessions_response import SearchSessionsResponse
from ...types.session import Session as RawSession
from ...types.session_type import SessionType


def _wrap_draft_sessions_pager(
raw_pager: SyncPager[RawDraftSession, ListDraftSessionsResponse],
client: TrueFoundryGateway,
) -> SyncPager[AgentDraftSession, ListDraftSessionsResponse]:
wrapped_items = [AgentDraftSession(d, client) for d in (raw_pager.items or [])]
# The list-owned and search endpoints both return pages of typing.Union[Session, DraftSession] items,
# differing only in the response envelope type, so a single generic pager wrapper covers both.
_SessionsResponseT = typing.TypeVar("_SessionsResponseT")

def get_next() -> typing.Optional[SyncPager[AgentDraftSession, ListDraftSessionsResponse]]:
if raw_pager.get_next is None:
return None
next_raw = raw_pager.get_next()
if next_raw is None:
return None
return _wrap_draft_sessions_pager(next_raw, 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,
)


def _wrap_owned_session(
raw: ListOwnedSessionsResponseDataItem,
def _wrap_raw_session(
raw: typing.Union[RawSession, RawDraftSession],
client: TrueFoundryGateway,
) -> typing.Union[AgentSession, AgentDraftSession]:
# Dispatch a raw owned-session union member into its enriched wrapper, keyed off the `type` discriminant.
# Dispatch a raw session union member into its enriched wrapper, keyed off the `type` discriminant.
if raw.type == "session/draft":
return AgentDraftSession(raw, client)
if raw.type == "session":
return AgentSession(raw, client)
raise ValueError(f"Unknown owned session type: {raw.type!r}")
raise ValueError(f"Unknown session type: {raw.type!r}")


def _wrap_owned_sessions_pager(
raw_pager: SyncPager[ListOwnedSessionsResponseDataItem, ListOwnedSessionsResponse],
def _wrap_sessions_pager(
raw_pager: SyncPager[typing.Union[RawSession, RawDraftSession], _SessionsResponseT],
client: TrueFoundryGateway,
) -> SyncPager[typing.Union[AgentSession, AgentDraftSession], ListOwnedSessionsResponse]:
wrapped_items = [_wrap_owned_session(s, client) for s in (raw_pager.items or [])]
) -> SyncPager[typing.Union[AgentSession, AgentDraftSession], _SessionsResponseT]:
wrapped_items = [_wrap_raw_session(s, client) for s in (raw_pager.items or [])]

def get_next() -> typing.Optional[
SyncPager[typing.Union[AgentSession, AgentDraftSession], ListOwnedSessionsResponse]
]:
def get_next() -> typing.Optional[SyncPager[typing.Union[AgentSession, AgentDraftSession], _SessionsResponseT]]:
if raw_pager.get_next is None:
return None
next_raw = raw_pager.get_next()
if next_raw is None:
return None
return _wrap_owned_sessions_pager(next_raw, client)
return _wrap_sessions_pager(next_raw, client)

return SyncPager(
get_next=get_next if raw_pager.has_next else None,
Expand All @@ -80,55 +65,77 @@ def get_next() -> typing.Optional[
)


async def _async_wrap_draft_sessions_pager(
raw_pager: AsyncPager[RawDraftSession, ListDraftSessionsResponse],
client: AsyncTrueFoundryGateway,
) -> AsyncPager[AsyncAgentDraftSession, ListDraftSessionsResponse]:
wrapped_items = [AsyncAgentDraftSession(d, client) for d in (raw_pager.items or [])]
def _wrap_draft_sessions_pager(
raw_pager: SyncPager[RawDraftSession, ListDraftSessionsResponse],
client: TrueFoundryGateway,
) -> SyncPager[AgentDraftSession, ListDraftSessionsResponse]:
wrapped_items = [AgentDraftSession(d, client) for d in (raw_pager.items or [])]

async def get_next() -> typing.Optional[AsyncPager[AsyncAgentDraftSession, ListDraftSessionsResponse]]:
def get_next() -> typing.Optional[SyncPager[AgentDraftSession, ListDraftSessionsResponse]]:
if raw_pager.get_next is None:
return None
next_raw = await raw_pager.get_next()
next_raw = raw_pager.get_next()
if next_raw is None:
return None
return await _async_wrap_draft_sessions_pager(next_raw, client)
return _wrap_draft_sessions_pager(next_raw, client)

return AsyncPager(
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,
)


def _async_wrap_owned_session(
raw: ListOwnedSessionsResponseDataItem,
def _async_wrap_raw_session(
raw: typing.Union[RawSession, RawDraftSession],
client: AsyncTrueFoundryGateway,
) -> typing.Union[AsyncAgentSession, AsyncAgentDraftSession]:
# Dispatch a raw owned-session union member into its enriched wrapper, keyed off the `type` discriminant.
# Dispatch a raw session union member into its enriched wrapper, keyed off the `type` discriminant.
if raw.type == "session/draft":
return AsyncAgentDraftSession(raw, client)
if raw.type == "session":
return AsyncAgentSession(raw, client)
raise ValueError(f"Unknown owned session type: {raw.type!r}")
raise ValueError(f"Unknown session type: {raw.type!r}")


async def _async_wrap_owned_sessions_pager(
raw_pager: AsyncPager[ListOwnedSessionsResponseDataItem, ListOwnedSessionsResponse],
async def _async_wrap_sessions_pager(
raw_pager: AsyncPager[typing.Union[RawSession, RawDraftSession], _SessionsResponseT],
client: AsyncTrueFoundryGateway,
) -> AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], ListOwnedSessionsResponse]:
wrapped_items = [_async_wrap_owned_session(s, client) for s in (raw_pager.items or [])]
) -> AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], _SessionsResponseT]:
wrapped_items = [_async_wrap_raw_session(s, client) for s in (raw_pager.items or [])]

async def get_next() -> typing.Optional[
AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], ListOwnedSessionsResponse]
AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], _SessionsResponseT]
]:
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_owned_sessions_pager(next_raw, client)
return await _async_wrap_sessions_pager(next_raw, 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,
)


async def _async_wrap_draft_sessions_pager(
raw_pager: AsyncPager[RawDraftSession, ListDraftSessionsResponse],
client: AsyncTrueFoundryGateway,
) -> AsyncPager[AsyncAgentDraftSession, ListDraftSessionsResponse]:
wrapped_items = [AsyncAgentDraftSession(d, client) for d in (raw_pager.items or [])]

async def get_next() -> typing.Optional[AsyncPager[AsyncAgentDraftSession, ListDraftSessionsResponse]]:
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_draft_sessions_pager(next_raw, client)

return AsyncPager(
get_next=get_next if raw_pager.has_next else None,
Expand Down Expand Up @@ -317,7 +324,72 @@ def list_owned_sessions(
end_timestamp=end_timestamp,
request_options=request_options,
)
return _wrap_owned_sessions_pager(raw_pager, self._client)
return _wrap_sessions_pager(raw_pager, self._client)

def search_sessions(
self,
*,
agent_name: typing.Optional[str] = None,
created_by_subject_id: typing.Optional[str] = None,
created_by_subject_type: typing.Optional[CreatedBySubjectType] = None,
session_type: typing.Optional[SessionType] = None,
session_id: typing.Optional[str] = None,
limit: typing.Optional[int] = 10,
order: typing.Optional[SearchSessionsOrder] = None,
page_token: typing.Optional[str] = None,
start_timestamp: typing.Optional[str] = None,
end_timestamp: typing.Optional[str] = None,
request_options: typing.Optional[RequestOptions] = None,
) -> SyncPager[typing.Union[AgentSession, AgentDraftSession], SearchSessionsResponse]:
"""
Search sessions visible to the caller across agents, spanning both saved sessions and drafts
(newest first by default). Tenant admins see all tenant sessions; agent managers see sessions
on agents they manage plus their own; other callers see only their own.

Parameters
----------
agent_name : typing.Optional[str]
Filter to sessions linked to this saved agent.
created_by_subject_id : typing.Optional[str]
Filter to sessions created by this subject id.
created_by_subject_type : typing.Optional[CreatedBySubjectType]
Optional subject type used with created_by_subject_id.
session_type : typing.Optional[SessionType]
Filter by session type. Omit to include both saved sessions and drafts.
session_id : typing.Optional[str]
Filter to a specific session id.
limit : typing.Optional[int]
Page size. Default 10.
order : typing.Optional[SearchSessionsOrder]
Sort by creation time. Default ``desc``.
page_token : typing.Optional[str]
Token from the previous response ``next_page_token``.
start_timestamp : typing.Optional[str]
Inclusive lower bound on ``created_at`` (ISO-8601).
end_timestamp : typing.Optional[str]
Inclusive upper bound on ``created_at`` (ISO-8601).
request_options : typing.Optional[RequestOptions]
Overrides client timeout, retries, headers, and stream reconnect.

Returns
-------
SyncPager[typing.Union[AgentSession, AgentDraftSession], SearchSessionsResponse]
Paginated matching sessions.
"""
raw_pager = self._client.agents.private.search_sessions(
agent_name=agent_name,
created_by_subject_id=created_by_subject_id,
created_by_subject_type=created_by_subject_type,
session_type=session_type,
session_id=session_id,
limit=limit,
order=order,
page_token=page_token,
start_timestamp=start_timestamp,
end_timestamp=end_timestamp,
request_options=request_options,
)
return _wrap_sessions_pager(raw_pager, self._client)

def download_sandbox_file(
self,
Expand All @@ -343,9 +415,7 @@ def download_sandbox_file(
typing.Iterator[bytes]
The downloaded sandbox file content, chunked.
"""
return self._client.agents.private.download_sandbox_file(
sandbox_id, path=path, request_options=request_options
)
return self._client.agents.private.download_sandbox_file(sandbox_id, path=path, request_options=request_options)


class AsyncPrivateAgentSessionClient:
Expand Down Expand Up @@ -529,7 +599,72 @@ async def list_owned_sessions(
end_timestamp=end_timestamp,
request_options=request_options,
)
return await _async_wrap_owned_sessions_pager(raw_pager, self._client)
return await _async_wrap_sessions_pager(raw_pager, self._client)

async def search_sessions(
self,
*,
agent_name: typing.Optional[str] = None,
created_by_subject_id: typing.Optional[str] = None,
created_by_subject_type: typing.Optional[CreatedBySubjectType] = None,
session_type: typing.Optional[SessionType] = None,
session_id: typing.Optional[str] = None,
limit: typing.Optional[int] = 10,
order: typing.Optional[SearchSessionsOrder] = None,
page_token: typing.Optional[str] = None,
start_timestamp: typing.Optional[str] = None,
end_timestamp: typing.Optional[str] = None,
request_options: typing.Optional[RequestOptions] = None,
) -> AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], SearchSessionsResponse]:
"""
Search sessions visible to the caller across agents, spanning both saved sessions and drafts
(newest first by default). Tenant admins see all tenant sessions; agent managers see sessions
on agents they manage plus their own; other callers see only their own.

Parameters
----------
agent_name : typing.Optional[str]
Filter to sessions linked to this saved agent.
created_by_subject_id : typing.Optional[str]
Filter to sessions created by this subject id.
created_by_subject_type : typing.Optional[CreatedBySubjectType]
Optional subject type used with created_by_subject_id.
session_type : typing.Optional[SessionType]
Filter by session type. Omit to include both saved sessions and drafts.
session_id : typing.Optional[str]
Filter to a specific session id.
limit : typing.Optional[int]
Page size. Default 10.
order : typing.Optional[SearchSessionsOrder]
Sort by creation time. Default ``desc``.
page_token : typing.Optional[str]
Token from the previous response ``next_page_token``.
start_timestamp : typing.Optional[str]
Inclusive lower bound on ``created_at`` (ISO-8601).
end_timestamp : typing.Optional[str]
Inclusive upper bound on ``created_at`` (ISO-8601).
request_options : typing.Optional[RequestOptions]
Overrides client timeout, retries, headers, and stream reconnect.

Returns
-------
AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], SearchSessionsResponse]
Paginated matching sessions.
"""
raw_pager = await self._client.agents.private.search_sessions(
agent_name=agent_name,
created_by_subject_id=created_by_subject_id,
created_by_subject_type=created_by_subject_type,
session_type=session_type,
session_id=session_id,
limit=limit,
order=order,
page_token=page_token,
start_timestamp=start_timestamp,
end_timestamp=end_timestamp,
request_options=request_options,
)
return await _async_wrap_sessions_pager(raw_pager, self._client)

async def download_sandbox_file(
self,
Expand Down
Loading