From 7a6c7455e041493f40928cd22fe4e751b05b5deb Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Tue, 29 Sep 2026 18:00:30 -0700 Subject: [PATCH 01/18] Add hand-written baseten.sandbox surface over the generated clients Sync and async SandboxClient for the control plane: create, get, get_info, list (cursor-paginated), delete. Sandbox/AsyncSandbox for the execution plane with fs (read with retry, write) and process (exec) subsystems; drives, codegen, system, and network raise NotImplementedError until wrapped. API keys are exchanged for a sandbox token at POST /v1/token, cached until shortly before expiry, and replaced on 401 TOKEN_REVOKED by re-sending the request. One shared transport pools connections across both planes; HTTP/2 is best-effort via the optional h2 package. Records convert to curated types (SandboxInfo, SandboxProcessInfo) with open string statuses. Creation sends only the fields the caller set, so server-side defaults apply. Co-Authored-By: Claude Code --- baseten/sandbox/__init__.py | 48 +++ baseten/sandbox/_auth.py | 240 +++++++++++++ baseten/sandbox/_client.py | 626 ++++++++++++++++++++++++++++++++++ baseten/sandbox/_errors.py | 88 +++++ baseten/sandbox/_info.py | 184 ++++++++++ baseten/sandbox/_retry.py | 128 +++++++ baseten/sandbox/_sandbox.py | 498 +++++++++++++++++++++++++++ scripts/docgen/conf.py | 6 + scripts/docgen/index.rst | 1 + tests/sandbox/test_surface.py | 469 +++++++++++++++++++++++++ 10 files changed, 2288 insertions(+) create mode 100644 baseten/sandbox/__init__.py create mode 100644 baseten/sandbox/_auth.py create mode 100644 baseten/sandbox/_client.py create mode 100644 baseten/sandbox/_errors.py create mode 100644 baseten/sandbox/_info.py create mode 100644 baseten/sandbox/_retry.py create mode 100644 baseten/sandbox/_sandbox.py create mode 100644 tests/sandbox/test_surface.py diff --git a/baseten/sandbox/__init__.py b/baseten/sandbox/__init__.py new file mode 100644 index 0000000..1aeb477 --- /dev/null +++ b/baseten/sandbox/__init__.py @@ -0,0 +1,48 @@ +"""Baseten sandbox SDK. + +Use :class:`SandboxClient` or :class:`AsyncSandboxClient` to create, find, +and delete sandboxes, and the :class:`Sandbox` or :class:`AsyncSandbox` they +hand out to work in one. +""" + +from baseten.sandbox._auth import AsyncSandboxTokenProvider, SandboxTokenProvider +from baseten.sandbox._client import ( + AsyncSandboxClient, + SandboxClient, + SandboxClientOptions, +) +from baseten.sandbox._errors import SandboxApiError, SandboxGatewayError +from baseten.sandbox._info import ( + SandboxEnvValue, + SandboxInfo, + SandboxProcessInfo, +) +from baseten.sandbox._retry import SandboxRetryOptions +from baseten.sandbox._sandbox import ( + AsyncSandbox, + AsyncSandboxFileSystem, + AsyncSandboxProcess, + Sandbox, + SandboxFileSystem, + SandboxProcess, +) + +__all__ = [ + "AsyncSandbox", + "AsyncSandboxClient", + "AsyncSandboxFileSystem", + "AsyncSandboxProcess", + "AsyncSandboxTokenProvider", + "Sandbox", + "SandboxApiError", + "SandboxClient", + "SandboxClientOptions", + "SandboxEnvValue", + "SandboxFileSystem", + "SandboxGatewayError", + "SandboxInfo", + "SandboxProcess", + "SandboxProcessInfo", + "SandboxRetryOptions", + "SandboxTokenProvider", +] diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py new file mode 100644 index 0000000..ffe907d --- /dev/null +++ b/baseten/sandbox/_auth.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +import asyncio +import threading +from collections.abc import Awaitable, Callable +from datetime import UTC, datetime, timedelta +from typing import Protocol + +import httpx + +# Refresh this long before the token's stated expiry, to allow for clock skew +# between this machine and the server, and for time spent in flight. +_TOKEN_EXPIRY_LEEWAY = timedelta(seconds=60) + +# A few retries cover a new token being invalidated in the same event as the +# old one. Past that, the failure is raised, since tokens that keep getting +# rejected point at something a new token cannot fix. +_TOKEN_INVALIDATION_MAX_RETRIES = 2 + +# The control plane and sandboxes send this header when they reject a request's +# token as revoked, which can happen before the token's stated expiry. The +# request was rejected before any work was done, so sending it again with a +# new token is safe even when it has side effects. +_TOKEN_REVOKED_CODE = "TOKEN_REVOKED" + + +class SandboxTokenProvider(Protocol): + """Returns the bearer token to send on a request. Called for every request.""" + + def __call__(self) -> str: ... + + +class AsyncSandboxTokenProvider(Protocol): + """Returns the bearer token to send on a request, or an awaitable of one.""" + + def __call__(self) -> str | Awaitable[str]: ... + + +def resolve_http2_flag(http2: bool | None) -> bool: + """Resolve the HTTP/2 option to a concrete flag for httpx. + + ``None`` enables HTTP/2 when the optional ``h2`` package is importable; + ``True`` requires it and fails otherwise; ``False`` never uses it. + """ + if http2 is False: + return False + try: + import h2 # ty: ignore[unresolved-import] # noqa: F401 + except ImportError: + if http2 is True: + raise ValueError( + "http2=True requires the h2 package; install baseten with the" + " http2 extra (pip install baseten[http2])" + ) + return False + return True + + +class _TokenCache: + """A cached minted token, guarded so concurrent callers share one mint.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._token: str | None = None + self._expires_at: datetime | None = None + + def get(self) -> tuple[str, datetime] | None: + with self._lock: + if self._token is None or self._expires_at is None: + return None + if datetime.now(UTC) >= self._expires_at - _TOKEN_EXPIRY_LEEWAY: + self._token = None + self._expires_at = None + return None + return self._token, self._expires_at + + def put(self, token: str, expires_at: datetime) -> None: + with self._lock: + self._token = token + self._expires_at = expires_at + + def drop(self, token: str) -> None: + """Drop a token the server rejected, so the next request gets a new one. + + Compared by value, so a rejection of an older token never discards a + newer one another request already minted. + """ + with self._lock: + if self._token == token: + self._token = None + self._expires_at = None + + +class SyncTokenSource: + """Where each request's bearer token comes from: a caller's token provider, + or a token minted from an API key, cached until shortly before it expires. + + The API key is exchanged for a token because the sandbox APIs do not + accept API keys directly. Should that change, only this class changes. + """ + + def __init__( + self, + *, + api_key: str, + token_provider: SandboxTokenProvider | None, + mint: Callable[[], tuple[str, datetime]] | None, + ) -> None: + if token_provider is not None and api_key != "": + raise ValueError("api_key must be empty when token_provider is set") + self._provider = token_provider + self._mint = mint + self._cache = _TokenCache() + + def token(self) -> str | None: + """The token for a request, or None to send no Authorization. + + None is the advanced opt-out of an empty API key with no provider. + """ + if self._provider is not None: + return self._provider() + if self._mint is None: + return None + cached = self._cache.get() + if cached is not None: + return cached[0] + token, expires_at = self._mint() + self._cache.put(token, expires_at) + return token + + def invalidate(self, token: str) -> None: + self._cache.drop(token) + + +class AsyncTokenSource: + """Async variant of :class:`SyncTokenSource`. + + The token provider may return a string or an awaitable of one. The + asyncio lock makes concurrent callers share one mint. + """ + + def __init__( + self, + *, + api_key: str, + token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None, + mint: Callable[[], Awaitable[tuple[str, datetime]]] | None, + ) -> None: + if token_provider is not None and api_key != "": + raise ValueError("api_key must be empty when token_provider is set") + self._provider = token_provider + self._mint = mint + self._lock = asyncio.Lock() + self._token: str | None = None + self._expires_at: datetime | None = None + + async def token(self) -> str | None: + if self._provider is not None: + provided = self._provider() + if isinstance(provided, str): + return provided + return await provided + if self._mint is None: + return None + async with self._lock: + if ( + self._token is not None + and self._expires_at is not None + and datetime.now(UTC) < self._expires_at - _TOKEN_EXPIRY_LEEWAY + ): + return self._token + token, expires_at = await self._mint() + self._token = token + self._expires_at = expires_at + return token + + async def invalidate(self, token: str) -> None: + async with self._lock: + if self._token == token: + self._token = None + self._expires_at = None + + +class _SyncAuthTransport(httpx.BaseTransport): + """Adds a bearer token from the source to every request. + + A request whose token was revoked is re-sent with a fresh one, up to + :data:`_TOKEN_INVALIDATION_MAX_RETRIES` times. The request body is + buffered first so a re-send replays it. + """ + + def __init__( + self, inner: httpx.BaseTransport, token_source: SyncTokenSource + ) -> None: + self._inner = inner + self._token_source = token_source + + def handle_request(self, request: httpx.Request) -> httpx.Response: + request.read() + for _ in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): + token = self._token_source.token() + if token is not None: + request.headers["Authorization"] = f"Bearer {token}" + response = self._inner.handle_request(request) + if ( + token is None + or response.status_code != 401 + or response.headers.get("x-blaxel-error-code") != _TOKEN_REVOKED_CODE + ): + return response + self._token_source.invalidate(token) + response.close() + raise AssertionError # pragma: no cover - loop returns first + + +class _AsyncAuthTransport(httpx.AsyncBaseTransport): + """Async variant of :class:`_SyncAuthTransport`.""" + + def __init__( + self, inner: httpx.AsyncBaseTransport, token_source: AsyncTokenSource + ) -> None: + self._inner = inner + self._token_source = token_source + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + request.read() + for _ in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): + token = await self._token_source.token() + if token is not None: + request.headers["Authorization"] = f"Bearer {token}" + response = await self._inner.handle_async_request(request) + if ( + token is None + or response.status_code != 401 + or response.headers.get("x-blaxel-error-code") != _TOKEN_REVOKED_CODE + ): + return response + await self._token_source.invalidate(token) + await response.aclose() + raise AssertionError # pragma: no cover - loop returns first diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py new file mode 100644 index 0000000..b0b18ae --- /dev/null +++ b/baseten/sandbox/_client.py @@ -0,0 +1,626 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Self + +import httpx + +import baseten.client.managementapi +from baseten.client._user_agent import with_user_agent +from baseten.sandbox._auth import ( + AsyncSandboxTokenProvider, + AsyncTokenSource, + SandboxTokenProvider, + SyncTokenSource, + _AsyncAuthTransport, + _SyncAuthTransport, + resolve_http2_flag, +) +from baseten.sandbox._errors import to_sandbox_api_error +from baseten.sandbox._info import ( + SandboxEnvValue, + SandboxInfo, + sandbox_envs_to_api, + sandbox_info_from_api, +) +from baseten.sandbox._retry import SandboxRetryOptions +from baseten.sandbox._sandbox import AsyncSandbox, Sandbox + +DEFAULT_MANAGEMENT_BASE_URL = "https://api.baseten.co" + +# Only connect time is bounded by default: a sandbox call can run for as long +# as the operation needs, and callers pass a stricter `timeout` when wanted. +DEFAULT_TIMEOUT = httpx.Timeout(None, connect=10.0) + + +@dataclass(frozen=True) +class SandboxClientOptions: + """Options for :class:`SandboxClient` and :class:`AsyncSandboxClient`. + + Obtain via :attr:`SandboxClient.options` to inspect the values a client + was constructed with. + """ + + api_key: str + """API key for authentication. Empty, together with *token_provider*, is + the advanced opt-out of API key authentication.""" + + token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None = None + """Returns the bearer token for each request instead of one minted from + the API key, which must then be empty.""" + + team_id: str | None = None + """Team to act in. Unset uses the caller's only team, and is an error for + callers in more than one team.""" + + base_url_override: str | None = None + """Explicit management API base URL, or None to use the default.""" + + headers: Mapping[str, str] | None = None + """Additional headers to send on every request.""" + + http2: bool | None = None + """Use HTTP/2 when the optional h2 package allows. None enables it when + h2 is importable, True requires it, False never uses it.""" + + timeout: httpx.Timeout | None = None + """HTTP timeouts. Defaults to bounding only the connect time.""" + + retries: SandboxRetryOptions | None = None + """Retry budgets for transient failures.""" + + transport_override: httpx.BaseTransport | httpx.AsyncBaseTransport | None = None + """HTTP transport to send every request through, both planes. For tests + with a mock transport.""" + + +def _management_base_url(options: SandboxClientOptions) -> str: + return options.base_url_override or DEFAULT_MANAGEMENT_BASE_URL + + +def _request_headers( + options: SandboxClientOptions, authorization: str | None +) -> dict[str, str]: + headers: dict[str, str] = {**(options.headers or {})} + if authorization is not None: + headers["Authorization"] = authorization + return with_user_agent(headers) + + +def _sandbox_create_request( + *, + name: str | None, + image: str | None, + memory: int | None, + region: str | None, + envs: list[baseten.client.managementapi.SandboxEnv] | None, + labels: dict[str, str] | None, + display_name: str | None, + external_id: str | None, +) -> baseten.client.managementapi.CreateSandboxRequest: + """Build the create body from only the fields the caller set. + + The generated client keeps explicitly passed None values, since null can + mean "clear" on updates. Creation must send nothing instead, so the + server-side defaults apply. + """ + create_fields: dict[str, Any] = {"name": name} + if image is not None: + create_fields["image"] = image + if memory is not None: + create_fields["memory"] = memory + if region is not None: + create_fields["region"] = region + if envs is not None: + create_fields["envs"] = envs + if labels is not None: + create_fields["labels"] = labels + if display_name is not None: + create_fields["display_name"] = display_name + if external_id is not None: + create_fields["external_id"] = external_id + return baseten.client.managementapi.CreateSandboxRequest(**create_fields) + + +class SandboxClient: + """Synchronous client for Baseten sandboxes. + + Creates, finds, and deletes sandboxes, and hands out :class:`Sandbox` + objects to work in one. Closing the client closes the sandboxes' shared + connections. Can be used as a context manager. + """ + + def __init__( + self, + *, + api_key: str, + token_provider: SandboxTokenProvider | None = None, + team_id: str | None = None, + base_url_override: str | None = None, + headers: Mapping[str, str] | None = None, + http2: bool | None = None, + timeout: httpx.Timeout | None = None, + retries: SandboxRetryOptions | None = None, + transport_override: httpx.BaseTransport | None = None, + ) -> None: + """Create a new synchronous sandbox client. + + Every argument matches the :class:`SandboxClientOptions` field of the + same name. + """ + self._options = SandboxClientOptions( + api_key=api_key, + token_provider=token_provider, + team_id=team_id, + base_url_override=base_url_override, + headers=headers, + http2=http2, + timeout=timeout, + retries=retries, + transport_override=transport_override, + ) + if transport_override is None: + self._pool: httpx.BaseTransport = httpx.HTTPTransport( + http2=resolve_http2_flag(self._options.http2) + ) + elif isinstance(transport_override, httpx.BaseTransport): + # MockTransport implements both interfaces, so an async-only + # transport (not a BaseTransport) is what must fall through. + self._pool = transport_override + else: + raise TypeError( + "SandboxClient needs a sync transport in transport_override" + ) + self._timeout = self._options.timeout or DEFAULT_TIMEOUT + # The mint client authenticates with the API key directly; every + # other client goes through the auth transport below. + self._mint_http_client = httpx.Client( + transport=self._pool, + base_url=_management_base_url(self._options), + headers=_request_headers( + self._options, + f"Bearer {api_key}" if api_key != "" else None, + ), + ) + self._token_source = SyncTokenSource( + api_key=api_key, + token_provider=token_provider, + mint=self._mint if api_key != "" else None, + ) + self._auth_transport = _SyncAuthTransport(self._pool, self._token_source) + self._http_client = httpx.Client( + transport=self._auth_transport, + base_url=_management_base_url(self._options), + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._api = baseten.client.managementapi.ApiClient(self._http_client) + self._sandbox_http_clients: list[httpx.Client] = [] + + def _mint(self) -> tuple[str, datetime]: + # The mint client authenticates with the API key, not with a token, + # so this does not go back through the auth transport. + minted = baseten.client.managementapi.ApiClient( + self._mint_http_client + ).post_token( + request=baseten.client.managementapi.CreateTokenRequest( + scopes=[baseten.client.managementapi.TokenScope.sandboxes] + ) + ) + return minted.token, minted.expires_at + + @property + def options(self) -> SandboxClientOptions: + """The options this client was constructed with.""" + return self._options + + @property + def api(self) -> baseten.client.managementapi.ApiClient: + """The generated management API client, authenticated for sandboxes. + + For operations this SDK does not wrap. The generated API surface is + not covered by stability guarantees and may change between versions. + """ + return self._api + + def create( + self, + *, + name: str | None = None, + image: str | None = None, + memory: int | None = None, + region: str | None = None, + envs: dict[str, SandboxEnvValue] | None = None, + labels: dict[str, str] | None = None, + display_name: str | None = None, + external_id: str | None = None, + ) -> Sandbox: + """Create a sandbox. + + Defaults for image and memory live server-side, so nothing is sent + for what the caller leaves unset. + """ + try: + sandbox = self._api.create_sandbox( + params=baseten.client.managementapi.CreateSandboxParams( + team_id=self._options.team_id + ), + request=_sandbox_create_request( + name=name, + image=image, + memory=memory, + region=region, + envs=sandbox_envs_to_api(envs) if envs is not None else None, + labels=labels, + display_name=display_name, + external_id=external_id, + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + info = sandbox_info_from_api(sandbox) + return self._sandbox_for(info, info) + + def get_info(self, name: str) -> SandboxInfo: + """Get a sandbox's current record.""" + try: + sandbox = self._api.get_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.GetSandboxParams( + team_id=self._options.team_id + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return sandbox_info_from_api(sandbox) + + def get(self, sandbox: str | SandboxInfo) -> Sandbox: + """Get a :class:`Sandbox` to work in. + + Fetches the record when given a name; none is fetched when the + record is already in hand. + """ + info = self.get_info(sandbox) if isinstance(sandbox, str) else sandbox + return self._sandbox_for(info) + + def list( + self, + *, + query: str | None = None, + statuses: Sequence[str] | None = None, + external_id: str | None = None, + page_size: int | None = None, + ) -> Iterator[SandboxInfo]: + """List sandboxes, fetching further pages as iteration reaches them.""" + seen_cursors: set[str] = set() + cursor: str | None = None + while True: + try: + response = self._api.list_sandboxes( + params=baseten.client.managementapi.ListSandboxesParams( + team_id=self._options.team_id, + cursor=cursor, + q=query, + status=[ + baseten.client.managementapi.StatusItem(root=status) + for status in statuses + ] + if statuses is not None + else None, + external_id=external_id, + **({"limit": page_size} if page_size is not None else {}), + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + yield from (sandbox_info_from_api(sb) for sb in response.items) + cursor = ( + response.pagination.cursor if response.pagination.has_more else None + ) + if cursor is None: + return + if cursor in seen_cursors: + raise ValueError("sandbox list returned a repeated cursor") + seen_cursors.add(cursor) + + def delete(self, name: str) -> SandboxInfo: + """Delete a sandbox. + + Deletion continues after this returns, so the returned record usually + shows the sandbox as still deleting. + """ + try: + sandbox = self._api.delete_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.DeleteSandboxParams( + team_id=self._options.team_id + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return sandbox_info_from_api(sandbox) + + def close(self) -> None: + """Close the client and the connections every sandbox shares.""" + for http_client in self._sandbox_http_clients: + http_client.close() + self._http_client.close() + self._mint_http_client.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def _sandbox_for( + self, info: SandboxInfo, creation_info: SandboxInfo | None = None + ) -> Sandbox: + if info.url is None: + raise ValueError(f"sandbox {info.name} has no URL yet") + http_client = httpx.Client( + transport=self._auth_transport, + base_url=info.url, + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._sandbox_http_clients.append(http_client) + return Sandbox._shared( + name=info.name, + url=info.url, + http_client=http_client, + retries=self._options.retries, + creation_info=creation_info, + ) + + +class AsyncSandboxClient: + """Asynchronous client for Baseten sandboxes. + + Creates, finds, and deletes sandboxes, and hands out + :class:`AsyncSandbox` objects to work in one. Closing the client closes + the sandboxes' shared connections. Can be used as an async context + manager. + """ + + def __init__( + self, + *, + api_key: str, + token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None = None, + team_id: str | None = None, + base_url_override: str | None = None, + headers: Mapping[str, str] | None = None, + http2: bool | None = None, + timeout: httpx.Timeout | None = None, + retries: SandboxRetryOptions | None = None, + transport_override: httpx.AsyncBaseTransport | None = None, + ) -> None: + """Create a new asynchronous sandbox client. + + The token provider may return a string or an awaitable of one. Every + other argument matches the :class:`SandboxClientOptions` field of + the same name. + """ + self._options = SandboxClientOptions( + api_key=api_key, + token_provider=token_provider, + team_id=team_id, + base_url_override=base_url_override, + headers=headers, + http2=http2, + timeout=timeout, + retries=retries, + transport_override=transport_override, + ) + if transport_override is None: + self._pool: httpx.AsyncBaseTransport = httpx.AsyncHTTPTransport( + http2=resolve_http2_flag(self._options.http2) + ) + elif isinstance(transport_override, httpx.AsyncBaseTransport): + self._pool = transport_override + else: + raise TypeError( + "AsyncSandboxClient needs an async transport in transport_override" + ) + self._timeout = self._options.timeout or DEFAULT_TIMEOUT + self._mint_http_client = httpx.AsyncClient( + transport=self._pool, + base_url=_management_base_url(self._options), + headers=_request_headers( + self._options, + f"Bearer {api_key}" if api_key != "" else None, + ), + ) + + async def mint() -> tuple[str, datetime]: + minted = await baseten.client.managementapi.AsyncApiClient( + self._mint_http_client + ).post_token( + request=baseten.client.managementapi.CreateTokenRequest( + scopes=[baseten.client.managementapi.TokenScope.sandboxes] + ) + ) + return minted.token, minted.expires_at + + self._token_source = AsyncTokenSource( + api_key=api_key, + token_provider=self._options.token_provider, # type: ignore[arg-type] + mint=mint if api_key != "" else None, + ) + self._auth_transport = _AsyncAuthTransport(self._pool, self._token_source) + self._http_client = httpx.AsyncClient( + transport=self._auth_transport, + base_url=_management_base_url(self._options), + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._api = baseten.client.managementapi.AsyncApiClient(self._http_client) + self._sandbox_http_clients: list[httpx.AsyncClient] = [] + + @property + def options(self) -> SandboxClientOptions: + """The options this client was constructed with.""" + return self._options + + @property + def api(self) -> baseten.client.managementapi.AsyncApiClient: + """The generated management API client, authenticated for sandboxes. + + For operations this SDK does not wrap. The generated API surface is + not covered by stability guarantees and may change between versions. + """ + return self._api + + async def create( + self, + *, + name: str | None = None, + image: str | None = None, + memory: int | None = None, + region: str | None = None, + envs: dict[str, SandboxEnvValue] | None = None, + labels: dict[str, str] | None = None, + display_name: str | None = None, + external_id: str | None = None, + ) -> AsyncSandbox: + """Create a sandbox. + + Defaults for image and memory live server-side, so nothing is sent + for what the caller leaves unset. + """ + try: + sandbox = await self._api.create_sandbox( + params=baseten.client.managementapi.CreateSandboxParams( + team_id=self._options.team_id + ), + request=_sandbox_create_request( + name=name, + image=image, + memory=memory, + region=region, + envs=sandbox_envs_to_api(envs) if envs is not None else None, + labels=labels, + display_name=display_name, + external_id=external_id, + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + info = sandbox_info_from_api(sandbox) + return await self._sandbox_for(info, info) + + async def get_info(self, name: str) -> SandboxInfo: + """Get a sandbox's current record.""" + try: + sandbox = await self._api.get_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.GetSandboxParams( + team_id=self._options.team_id + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return sandbox_info_from_api(sandbox) + + async def get(self, sandbox: str | SandboxInfo) -> AsyncSandbox: + """Get an :class:`AsyncSandbox` to work in. + + Fetches the record when given a name; none is fetched when the + record is already in hand. + """ + info = await self.get_info(sandbox) if isinstance(sandbox, str) else sandbox + return await self._sandbox_for(info) + + async def list( + self, + *, + query: str | None = None, + statuses: Sequence[str] | None = None, + external_id: str | None = None, + page_size: int | None = None, + ) -> AsyncIterator[SandboxInfo]: + """List sandboxes, fetching further pages as iteration reaches them.""" + seen_cursors: set[str] = set() + cursor: str | None = None + while True: + try: + response = await self._api.list_sandboxes( + params=baseten.client.managementapi.ListSandboxesParams( + team_id=self._options.team_id, + cursor=cursor, + q=query, + status=[ + baseten.client.managementapi.StatusItem(root=status) + for status in statuses + ] + if statuses is not None + else None, + external_id=external_id, + **({"limit": page_size} if page_size is not None else {}), + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + for sb in response.items: + yield sandbox_info_from_api(sb) + cursor = ( + response.pagination.cursor if response.pagination.has_more else None + ) + if cursor is None: + return + if cursor in seen_cursors: + raise ValueError("sandbox list returned a repeated cursor") + seen_cursors.add(cursor) + + async def delete(self, name: str) -> SandboxInfo: + """Delete a sandbox. + + Deletion continues after this returns, so the returned record usually + shows the sandbox as still deleting. + """ + try: + sandbox = await self._api.delete_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.DeleteSandboxParams( + team_id=self._options.team_id + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return sandbox_info_from_api(sandbox) + + async def close(self) -> None: + """Close the client and the connections every sandbox shares.""" + for http_client in self._sandbox_http_clients: + await http_client.aclose() + await self._http_client.aclose() + await self._mint_http_client.aclose() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *args: object) -> None: + await self.close() + + async def _sandbox_for( + self, info: SandboxInfo, creation_info: SandboxInfo | None = None + ) -> AsyncSandbox: + if info.url is None: + raise ValueError(f"sandbox {info.name} has no URL yet") + http_client = httpx.AsyncClient( + transport=self._auth_transport, + base_url=info.url, + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._sandbox_http_clients.append(http_client) + return AsyncSandbox._shared( + name=info.name, + url=info.url, + http_client=http_client, + retries=self._options.retries, + creation_info=creation_info, + ) diff --git a/baseten/sandbox/_errors.py b/baseten/sandbox/_errors.py new file mode 100644 index 0000000..64ac7b6 --- /dev/null +++ b/baseten/sandbox/_errors.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import json + +import baseten.client.managementapi +import baseten.client.sandboxapi + + +class SandboxApiError(Exception): + """An error response from the sandbox control plane or from a sandbox. + + The response body is parsed as JSON when possible, otherwise kept as + text. *code* holds the machine-readable error code when the response has + one; the control plane puts it in ``code`` or ``error``, a sandbox in + ``error``. + """ + + def __init__(self, status: int, body: object) -> None: + fields = body if isinstance(body, dict) else {} + body_message = fields.get("message") + body_error = fields.get("error") + body_code = fields.get("code") + code = ( + body_code + if isinstance(body_code, str) + else (body_error if isinstance(body_message, str) else None) + ) + description = None + if isinstance(body_message, str): + description = body_message + elif isinstance(body_error, str): + description = body_error + message = f"sandbox API error (HTTP {status})" + if isinstance(code, str): + message += f" {code}" + if description is not None: + message += f": {description}" + super().__init__(message) + self.status = status + self.code = code + self.details = fields.get("details") + self.body = body + + +class SandboxGatewayError(SandboxApiError): + """A 502, 503, or 504 from the edge in front of a sandbox. + + Not from the sandbox itself: usually transient, for example while a + sandbox wakes from standby or when a request outlasts the edge's timeout. + """ + + +_GATEWAY_STATUSES = frozenset({502, 503, 504}) + + +def to_sandbox_api_error(error: object, plane: str) -> Exception: + """Convert an error raised by a generated client, or return it unchanged. + + ``plane`` is ``"control"`` or ``"exec"``. Only a sandbox sits behind the + edge, so gateway statuses from the control plane stay ordinary errors. + """ + status = getattr(error, "status_code", None) + if not isinstance(status, int): + return error if isinstance(error, Exception) else Exception(str(error)) + if isinstance( + error, + ( + baseten.client.managementapi.ResponseError, + baseten.client.sandboxapi.ResponseError, + ), + ): + try: + body: object = json.loads(error.body) + except (json.JSONDecodeError, TypeError): + body = error.body + else: + # Generated typed errors are dataclasses of status_code plus the + # error's parsed model, keyed by the schema name. + model_fields = { + name: value for name, value in vars(error).items() if name != "status_code" + } + if len(model_fields) == 1: + body = next(iter(model_fields.values())) + else: # pragma: no cover - generated errors carry exactly one model + body = model_fields + if plane == "exec" and status in _GATEWAY_STATUSES: + return SandboxGatewayError(status, body) + return SandboxApiError(status, body) diff --git a/baseten/sandbox/_info.py b/baseten/sandbox/_info.py new file mode 100644 index 0000000..bf6f1e9 --- /dev/null +++ b/baseten/sandbox/_info.py @@ -0,0 +1,184 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime +from email.utils import parsedate_to_datetime + +import baseten.client.managementapi +import baseten.client.sandboxapi + +# Deployment status of a sandbox. Only DEPLOYED sandboxes run commands. +# Kept open: the server may add values, so do not treat this as exhaustive. +SANDBOX_STATUSES = ( + "DEPLOYING", + "DEPLOYED", + "FAILED", + "DEACTIVATING", + "DEACTIVATED", + "DELETING", + "TERMINATED", + "ARCHIVING", + "ARCHIVED", + "UNARCHIVING", + "BUILDING", + "UPLOADING", +) + +# Whether a deployed sandbox is running or idle in standby. Kept open. +SANDBOX_STATES = ("RUNNING", "STANDBY") + +# Status of a process in a sandbox. Kept open. +SANDBOX_PROCESS_STATUSES = ("running", "completed", "failed", "killed", "stopped") + +SandboxStatus = str +SandboxState = str +SandboxProcessStatus = str + + +@dataclass +class SandboxEnvValue: + """Value of an environment variable in a sandbox.""" + + value: str + secret: bool = False + + +@dataclass +class SandboxInfo: + """A sandbox as last reported by the control plane.""" + + name: str + """Unique name of the sandbox, assigned by the server when not given.""" + + status: SandboxStatus + enabled: bool + created_at: datetime + envs: dict[str, SandboxEnvValue] = field(default_factory=dict) + labels: dict[str, str] = field(default_factory=dict) + + url: str | None = None + """Base URL of the sandbox's execution API, once it has one.""" + + state: SandboxState | None = None + image: str | None = None + """Image reference, including its tag.""" + + memory: int | None = None + """Memory in megabytes, which also sets the CPU allocation.""" + + region: str | None = None + display_name: str | None = None + external_id: str | None = None + """Caller-owned identifier for external lookups.""" + + updated_at: datetime | None = None + created_by: str | None = None + updated_by: str | None = None + last_used_at: datetime | None = None + expires_in_seconds: int | None = None + """Seconds left before automatic deletion, when expiration is configured.""" + + +@dataclass +class SandboxProcessInfo: + """A process in a sandbox.""" + + pid: str + name: str + command: str + status: SandboxProcessStatus + exit_code: int + stdout: str + stderr: str + logs: str + """Standard output and standard error, interleaved.""" + + working_dir: str + started_at: datetime + completed_at: datetime | None = None + + +def sandbox_info_from_api(sandbox: baseten.client.managementapi.Sandbox) -> SandboxInfo: + """Convert the control plane's sandbox record.""" + return SandboxInfo( + name=sandbox.name, + url=str(sandbox.url) if sandbox.url is not None else None, + status=str(sandbox.status), + state=str(sandbox.state) if sandbox.state is not None else None, + image=sandbox.image, + memory=sandbox.memory, + region=sandbox.region, + enabled=sandbox.enabled, + envs=_envs_from_api(sandbox.envs), + labels=dict(sandbox.labels.root) if sandbox.labels is not None else {}, + display_name=sandbox.display_name, + external_id=sandbox.external_id, + created_at=sandbox.created_at, + updated_at=sandbox.updated_at, + created_by=sandbox.created_by, + updated_by=sandbox.updated_by, + last_used_at=sandbox.last_used_at, + expires_in_seconds=sandbox.expires_in, + ) + + +def sandbox_envs_to_api( + envs: dict[str, SandboxEnvValue], +) -> list[baseten.client.managementapi.SandboxEnv]: + """Convert environment variables to the control plane's form.""" + return [ + baseten.client.managementapi.SandboxEnv( + name=name, value=value.value, secret=value.secret + ) + for name, value in envs.items() + ] + + +def process_info_from_api( + process: baseten.client.sandboxapi.ProcessResponse, +) -> SandboxProcessInfo: + """Convert the execution plane's process record. + + The plane reports timestamps as strings in either ISO 8601 or RFC 1123 + form, and an empty completion time while the process runs. + """ + return SandboxProcessInfo( + pid=process.pid, + name=process.name, + command=process.command, + status=str(process.status), + exit_code=process.exitCode, + stdout=process.stdout, + stderr=process.stderr, + logs=process.logs, + working_dir=process.workingDir, + started_at=parse_exec_timestamp(process.startedAt), + completed_at=( + parse_exec_timestamp(process.completedAt) if process.completedAt else None + ), + ) + + +def parse_exec_timestamp(value: str) -> datetime: + """Parse an execution-plane timestamp, ISO 8601 or RFC 1123.""" + try: + return datetime.fromisoformat(value) + except ValueError: + pass + parsed = parsedate_to_datetime(value) + if parsed is None: + raise ValueError(f"unparseable execution-plane timestamp {value!r}") from None + return parsed + + +def _envs_from_api( + envs: list[baseten.client.managementapi.SandboxEnv] | None, +) -> dict[str, SandboxEnvValue]: + result: dict[str, SandboxEnvValue] = {} + for env in envs or []: + if env.name is None: + continue + result[env.name] = SandboxEnvValue( + value=env.value or "", secret=bool(env.secret) + ) + return result diff --git a/baseten/sandbox/_retry.py b/baseten/sandbox/_retry.py new file mode 100644 index 0000000..86db63d --- /dev/null +++ b/baseten/sandbox/_retry.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import asyncio +import random +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import TypeVar + +import httpx + +from baseten.sandbox._errors import SandboxGatewayError + +_T = TypeVar("_T") + +DEFAULT_READ_MAX_RETRIES = 5 +DEFAULT_GATEWAY_MAX_RETRIES = 2 +DEFAULT_UPLOAD_MAX_RETRIES = 3 + +_BACKOFF_BASE_SECONDS = 0.2 +_BACKOFF_MAX_SECONDS = 2.0 + +# httpx exceptions meaning the connection dropped before a complete response +# arrived. Only then is repeating an idempotent call safe. +_TRANSIENT_ERRORS = ( + httpx.ConnectError, + httpx.ConnectTimeout, + httpx.ReadError, + httpx.WriteError, + httpx.RemoteProtocolError, +) + + +@dataclass(frozen=True) +class SandboxRetryOptions: + """Retry budgets for transient failures. + + Retries only ever apply to operations that are safe to repeat, such as + reads. Operations with side effects, such as running a process or + creating a sandbox, are never retried. Each budget counts retries after + the first attempt, so 0 disables that kind of retry. + """ + + read_max_retries: int = DEFAULT_READ_MAX_RETRIES + """Retries of a read after a dropped or reset connection.""" + + gateway_max_retries: int = DEFAULT_GATEWAY_MAX_RETRIES + """Retries after the edge in front of a sandbox answers 502, 503, or 504. + + Each attempt can take as long as the edge's own timeout, so this is kept + small. + """ + + upload_max_retries: int = DEFAULT_UPLOAD_MAX_RETRIES + """Retries of a file upload after a dropped or reset connection.""" + + +def is_transient_reset_error(error: object) -> bool: + """Whether the error is a dropped or reset connection, not a response. + + A response from the server, even an error, means the request was handled; + only a transport failure makes repeating an idempotent call safe. + """ + return isinstance(error, _TRANSIENT_ERRORS) + + +def backoff_delay_seconds(attempt: int) -> float: + """Exponential backoff capped at a maximum, plus jitter. + + The jitter is up to one base delay, so many clients failing together do + not retry together. + """ + capped = min(_BACKOFF_BASE_SECONDS * 2 ** (attempt - 1), _BACKOFF_MAX_SECONDS) + return capped + random.uniform(0, _BACKOFF_BASE_SECONDS) + + +def retry_idempotent( + call: Callable[[], _T], + *, + max_retries: int, + gateway_max_retries: int, +) -> _T: + """Run an idempotent call, retrying dropped connections and edge errors. + + Must never wrap a call with side effects, since a dropped connection does + not tell whether the server acted on it. + """ + attempt = 0 + gateway_attempt = 0 + while True: + try: + return call() + except SandboxGatewayError: + gateway_attempt += 1 + if gateway_attempt > gateway_max_retries: + raise + delay = backoff_delay_seconds(gateway_attempt) + except Exception as error: + attempt += 1 + if attempt > max_retries or not is_transient_reset_error(error): + raise + delay = backoff_delay_seconds(attempt) + time.sleep(delay) + + +async def retry_idempotent_async( + call: Callable[[], Awaitable[_T]], + *, + max_retries: int, + gateway_max_retries: int, +) -> _T: + """Async variant of :func:`retry_idempotent`.""" + attempt = 0 + gateway_attempt = 0 + while True: + try: + return await call() + except SandboxGatewayError: + gateway_attempt += 1 + if gateway_attempt > gateway_max_retries: + raise + delay = backoff_delay_seconds(gateway_attempt) + except Exception as error: + attempt += 1 + if attempt > max_retries or not is_transient_reset_error(error): + raise + delay = backoff_delay_seconds(attempt) + await asyncio.sleep(delay) diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py new file mode 100644 index 0000000..fdf26a3 --- /dev/null +++ b/baseten/sandbox/_sandbox.py @@ -0,0 +1,498 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Self + +import httpx + +import baseten.client.sandboxapi +from baseten.client._user_agent import with_user_agent +from baseten.sandbox._auth import ( + AsyncTokenSource, + SandboxTokenProvider, + SyncTokenSource, + _AsyncAuthTransport, + _SyncAuthTransport, + resolve_http2_flag, +) +from baseten.sandbox._errors import to_sandbox_api_error +from baseten.sandbox._info import ( + SandboxInfo, + SandboxProcessInfo, + process_info_from_api, +) +from baseten.sandbox._retry import ( + SandboxRetryOptions, + retry_idempotent, + retry_idempotent_async, +) + +_NOT_IMPLEMENTED = ( + "Not implemented yet. Reach the operation through {} in the meantime." +) + + +class Sandbox: + """One sandbox, reached directly at its own URL. + + Get one from :meth:`baseten.sandbox.SandboxClient.create` or + :meth:`baseten.sandbox.SandboxClient.get`, or construct one directly when + the URL and a token provider are already known. A sandbox from a client + shares that client's connections and is closed with it; one constructed + directly owns its connections and :meth:`close` closes them. + """ + + def __init__( + self, + *, + name: str, + url: str, + token_provider: SandboxTokenProvider | None = None, + headers: Mapping[str, str] | None = None, + http2: bool | None = None, + timeout: httpx.Timeout | None = None, + retries: SandboxRetryOptions | None = None, + transport_override: httpx.BaseTransport | None = None, + ) -> None: + """Construct a standalone sandbox client. + + Args: + name: Name of the sandbox. + url: Base URL of the sandbox's execution API, as reported in + ``SandboxInfo.url``. + token_provider: Returns the bearer token for each request. When + unset, no Authorization is sent. + headers: Additional headers to send on every request. + http2: Use HTTP/2 when the optional h2 package allows. None + enables it when h2 is importable, True requires it. + timeout: HTTP timeouts. Defaults to bounding only connect time. + retries: Retry budgets for transient failures. + transport_override: HTTP transport for tests with a mock. + """ + self._name = name + self._url = url.rstrip("/") + self._retries = retries or SandboxRetryOptions() + self._pool = ( + transport_override + if transport_override is not None + else httpx.HTTPTransport(http2=resolve_http2_flag(http2)) + ) + self._token_source = SyncTokenSource( + api_key="", token_provider=token_provider, mint=None + ) + self._auth_transport = _SyncAuthTransport(self._pool, self._token_source) + request_headers = with_user_agent({**(headers or {})}) + self._http_client = httpx.Client( + transport=self._auth_transport, + base_url=self._url, + headers=request_headers, + timeout=timeout + if timeout is not None + else httpx.Timeout(None, connect=10.0), + ) + self._api = baseten.client.sandboxapi.ApiClient(self._http_client) + self._fs: SandboxFileSystem | None = None + self._process: SandboxProcess | None = None + self._info: SandboxInfo | None = None + self._owns_pool = True + + @classmethod + def _shared( + cls, + *, + name: str, + url: str, + http_client: httpx.Client, + retries: SandboxRetryOptions | None, + creation_info: SandboxInfo | None, + ) -> Sandbox: + sandbox = cls.__new__(cls) + sandbox._name = name + sandbox._url = url.rstrip("/") + sandbox._retries = retries or SandboxRetryOptions() + sandbox._http_client = http_client + sandbox._api = baseten.client.sandboxapi.ApiClient(http_client) + sandbox._fs = None + sandbox._process = None + sandbox._info = creation_info + sandbox._owns_pool = False + return sandbox + + @property + def name(self) -> str: + """Name of the sandbox.""" + return self._name + + @property + def url(self) -> str: + """Base URL of the sandbox's execution API.""" + return self._url + + @property + def info(self) -> SandboxInfo | None: + """The record as of creation, when this sandbox came from + :meth:`baseten.sandbox.SandboxClient.create`. None otherwise.""" + return self._info + + @property + def api(self) -> baseten.client.sandboxapi.ApiClient: + """The generated client for this sandbox's execution API. + + For operations this SDK does not wrap. The generated API surface is + not covered by stability guarantees and may change between versions. + """ + return self._api + + @property + def fs(self) -> SandboxFileSystem: + """Files and directories in the sandbox.""" + if self._fs is None: + self._fs = SandboxFileSystem(self._api, self._retries) + return self._fs + + @property + def process(self) -> SandboxProcess: + """Processes in the sandbox.""" + if self._process is None: + self._process = SandboxProcess(self._api) + return self._process + + @property + def drives(self) -> object: + """Drive mounts in the sandbox. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def codegen(self) -> object: + """Code editing helpers. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def system(self) -> object: + """Sandbox system operations. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def network(self) -> object: + """Network access to ports in the sandbox. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + def close(self) -> None: + """Close this sandbox's connections, when it owns them. + + A sandbox from a :class:`baseten.sandbox.SandboxClient` shares that + client's connections and is closed with it, so this does nothing. + """ + if self._owns_pool: + self._http_client.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + +class AsyncSandbox: + """One sandbox, reached directly at its own URL. Async variant. + + Get one from :meth:`baseten.sandbox.AsyncSandboxClient.create` or + :meth:`baseten.sandbox.AsyncSandboxClient.get`, or construct one directly + when the URL and a token provider are already known. + """ + + def __init__( + self, + *, + name: str, + url: str, + token_provider: SandboxTokenProvider | None = None, + headers: Mapping[str, str] | None = None, + http2: bool | None = None, + timeout: httpx.Timeout | None = None, + retries: SandboxRetryOptions | None = None, + transport_override: httpx.AsyncBaseTransport | None = None, + ) -> None: + """Construct a standalone async sandbox client. + + The token provider may return a string or an awaitable of one. Other + args match :class:`Sandbox`. + """ + self._name = name + self._url = url.rstrip("/") + self._retries = retries or SandboxRetryOptions() + self._pool = ( + transport_override + if transport_override is not None + else httpx.AsyncHTTPTransport(http2=resolve_http2_flag(http2)) + ) + self._token_source = AsyncTokenSource( + api_key="", + token_provider=token_provider, + mint=None, # type: ignore[arg-type] + ) + self._auth_transport = _AsyncAuthTransport(self._pool, self._token_source) + request_headers = with_user_agent({**(headers or {})}) + self._http_client = httpx.AsyncClient( + transport=self._auth_transport, + base_url=self._url, + headers=request_headers, + timeout=timeout + if timeout is not None + else httpx.Timeout(None, connect=10.0), + ) + self._api = baseten.client.sandboxapi.AsyncApiClient(self._http_client) + self._fs: AsyncSandboxFileSystem | None = None + self._process: AsyncSandboxProcess | None = None + self._info: SandboxInfo | None = None + self._owns_pool = True + + @classmethod + def _shared( + cls, + *, + name: str, + url: str, + http_client: httpx.AsyncClient, + retries: SandboxRetryOptions | None, + creation_info: SandboxInfo | None, + ) -> AsyncSandbox: + sandbox = cls.__new__(cls) + sandbox._name = name + sandbox._url = url.rstrip("/") + sandbox._retries = retries or SandboxRetryOptions() + sandbox._http_client = http_client + sandbox._api = baseten.client.sandboxapi.AsyncApiClient(http_client) + sandbox._fs = None + sandbox._process = None + sandbox._info = creation_info + sandbox._owns_pool = False + return sandbox + + @property + def name(self) -> str: + """Name of the sandbox.""" + return self._name + + @property + def url(self) -> str: + """Base URL of the sandbox's execution API.""" + return self._url + + @property + def info(self) -> SandboxInfo | None: + """The record as of creation, when this sandbox came from + :meth:`baseten.sandbox.AsyncSandboxClient.create`. None otherwise.""" + return self._info + + @property + def api(self) -> baseten.client.sandboxapi.AsyncApiClient: + """The generated client for this sandbox's execution API. + + For operations this SDK does not wrap. The generated API surface is + not covered by stability guarantees and may change between versions. + """ + return self._api + + @property + def fs(self) -> AsyncSandboxFileSystem: + """Files and directories in the sandbox.""" + if self._fs is None: + self._fs = AsyncSandboxFileSystem(self._api, self._retries) + return self._fs + + @property + def process(self) -> AsyncSandboxProcess: + """Processes in the sandbox.""" + if self._process is None: + self._process = AsyncSandboxProcess(self._api) + return self._process + + @property + def drives(self) -> object: + """Drive mounts in the sandbox. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def codegen(self) -> object: + """Code editing helpers. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def system(self) -> object: + """Sandbox system operations. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def network(self) -> object: + """Network access to ports in the sandbox. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + async def close(self) -> None: + """Close this sandbox's connections, when it owns them.""" + if self._owns_pool: + await self._http_client.aclose() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *args: object) -> None: + await self.close() + + +class SandboxFileSystem: + """Files and directories in a sandbox.""" + + def __init__( + self, api: baseten.client.sandboxapi.ApiClient, retries: SandboxRetryOptions + ) -> None: + self._api = api + self._retries = retries + + def read(self, path: str) -> str: + """Read a text file. + + Retries dropped connections and edge gateway errors, since a read is + safe to repeat. + """ + try: + result = retry_idempotent( + lambda: self._api.get_filesystem(path=path), + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + root = result.root + if isinstance(root, baseten.client.sandboxapi.FileWithContent): + return root.content + if isinstance(root, bytes): + return root.decode() + raise ValueError(f"{path} is a directory, not a file") + + def write(self, path: str, content: str) -> None: + """Write a text file, creating it or replacing its content.""" + try: + self._api.put_filesystem( + path=path, + request=baseten.client.sandboxapi.FileRequest(content=content), + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + +class AsyncSandboxFileSystem: + """Files and directories in a sandbox. Async variant.""" + + def __init__( + self, + api: baseten.client.sandboxapi.AsyncApiClient, + retries: SandboxRetryOptions, + ) -> None: + self._api = api + self._retries = retries + + async def read(self, path: str) -> str: + """Read a text file. + + Retries dropped connections and edge gateway errors, since a read is + safe to repeat. + """ + try: + result = await retry_idempotent_async( + lambda: self._api.get_filesystem(path=path), + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + root = result.root + if isinstance(root, baseten.client.sandboxapi.FileWithContent): + return root.content + if isinstance(root, bytes): + return root.decode() + raise ValueError(f"{path} is a directory, not a file") + + async def write(self, path: str, content: str) -> None: + """Write a text file, creating it or replacing its content.""" + try: + await self._api.put_filesystem( + path=path, + request=baseten.client.sandboxapi.FileRequest(content=content), + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + +class SandboxProcess: + """Processes in a sandbox.""" + + def __init__(self, api: baseten.client.sandboxapi.ApiClient) -> None: + self._api = api + + def exec( + self, + command: str, + *, + working_dir: str | None = None, + env: dict[str, str] | None = None, + name: str | None = None, + timeout_seconds: int | None = None, + wait_for_completion: bool | None = None, + ) -> SandboxProcessInfo: + """Start a command. + + Never retried, since a command may have side effects even when its + response is lost. + """ + try: + response = self._api.post_process( + request=baseten.client.sandboxapi.ProcessRequest( + command=command, + workingDir=working_dir, + env=env, + name=name, + timeout=timeout_seconds, + waitForCompletion=wait_for_completion, + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + return process_info_from_api(response) + + +class AsyncSandboxProcess: + """Processes in a sandbox. Async variant.""" + + def __init__(self, api: baseten.client.sandboxapi.AsyncApiClient) -> None: + self._api = api + + async def exec( + self, + command: str, + *, + working_dir: str | None = None, + env: dict[str, str] | None = None, + name: str | None = None, + timeout_seconds: int | None = None, + wait_for_completion: bool | None = None, + ) -> SandboxProcessInfo: + """Start a command. + + Never retried, since a command may have side effects even when its + response is lost. + """ + try: + response = await self._api.post_process( + request=baseten.client.sandboxapi.ProcessRequest( + command=command, + workingDir=working_dir, + env=env, + name=name, + timeout=timeout_seconds, + waitForCompletion=wait_for_completion, + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + return process_info_from_api(response) diff --git a/scripts/docgen/conf.py b/scripts/docgen/conf.py index 3888bd0..f6153c8 100644 --- a/scripts/docgen/conf.py +++ b/scripts/docgen/conf.py @@ -55,6 +55,12 @@ "baseten.client.inferenceapi": ["ApiClient", "AsyncApiClient"], "baseten.client.sandboxapi": ["ApiClient", "AsyncApiClient"], "baseten.client.modelconfig": ["ModelConfig"], + "baseten.sandbox": [ + "SandboxClient", + "AsyncSandboxClient", + "Sandbox", + "AsyncSandbox", + ], } diff --git a/scripts/docgen/index.rst b/scripts/docgen/index.rst index b4e7482..6e16b07 100644 --- a/scripts/docgen/index.rst +++ b/scripts/docgen/index.rst @@ -11,3 +11,4 @@ the source. For installation, usage, and examples, see the project on GitHub: :caption: API Reference baseten.client + baseten.sandbox diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py new file mode 100644 index 0000000..13ec95d --- /dev/null +++ b/tests/sandbox/test_surface.py @@ -0,0 +1,469 @@ +from __future__ import annotations + +import json +from typing import Any + +import httpx +import pytest + +from baseten.sandbox import ( + AsyncSandboxClient, + SandboxApiError, + SandboxClient, + SandboxGatewayError, + SandboxInfo, +) + +MANAGEMENT_URL = "https://api.baseten.co" +SANDBOX_URL = "https://sb-1.b10.co" +FAR_FUTURE = "2100-01-01T00:00:00Z" + + +def sandbox_record(name: str = "sb-1", status: str = "DEPLOYED") -> dict[str, Any]: + return { + "name": name, + "url": SANDBOX_URL, + "status": status, + "state": "RUNNING", + "created_at": "2026-09-29T00:00:00Z", + "enabled": True, + } + + +def process_record() -> dict[str, Any]: + return { + "command": "echo hi", + "completedAt": "Tue, 30 Sep 2026 00:00:01 GMT", + "exitCode": 0, + "logs": "hi\n", + "name": "echo", + "pid": "123", + "startedAt": "2026-09-30T00:00:00Z", + "status": "completed", + "stderr": "", + "stdout": "hi\n", + "workingDir": "/tmp", + } + + +def file_record() -> dict[str, Any]: + return { + "content": "file text", + "group": "root", + "lastModified": "2026-09-29T00:00:00Z", + "name": "file.txt", + "owner": "root", + "path": "/file.txt", + "permissions": "0644", + "size": 9, + } + + +class RoutingTransport: + """Mock transport routing by path, with scripted per-path responses. + + Each route maps path to either a response factory or a list consumed one + entry per matching request. + """ + + def __init__( + self, + routes: dict[str, Any], + mint_tokens: list[str] | None = None, + ) -> None: + self.routes = routes + self.mint_tokens = mint_tokens or ["tok-1"] + self.requests: list[tuple[str, str, str | None, bytes]] = [] + self.mint_count = 0 + + def _respond(self, request: httpx.Request) -> httpx.Response: + method = request.method + path = request.url.path + authorization = request.headers.get("authorization") + self.requests.append((method, path, authorization, request.content)) + + def json_response(payload: Any, status: int = 200) -> httpx.Response: + return httpx.Response( + status, json=payload, headers={"content-type": "application/json"} + ) + + if path == "/v1/token": + token = self.mint_tokens[min(self.mint_count, len(self.mint_tokens) - 1)] + self.mint_count += 1 + return json_response({"token": token, "expires_at": FAR_FUTURE}) + + route = self.routes.get(f"{method} {path}") + if route is None: + route = self.routes.get(path) + if route is None: + return httpx.Response(404, text=f"no route for {method} {path}") + if isinstance(route, list): + if not route: + return httpx.Response(500, text="route exhausted") + route = route.pop(0) + if isinstance(route, Exception): + raise route + if callable(route): + return route(request) + if isinstance(route, httpx.Response): + return route + return json_response(route) + + @property + def sync_transport(self) -> httpx.MockTransport: + return httpx.MockTransport(self._respond) + + @property + def async_transport(self) -> httpx.MockTransport: + async def handler(request: httpx.Request) -> httpx.Response: + return self._respond(request) + + return httpx.MockTransport(handler) + + def authorizations(self, path: str) -> list[str | None]: + return [auth for _, p, auth, _ in self.requests if p == path and auth] + + def count(self, method: str, path: str) -> int: + return sum(1 for m, p, _, _ in self.requests if m == method and p == path) + + +def make_sync_client(transport: RoutingTransport) -> SandboxClient: + return SandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.sync_transport, + ) + + +def created_response(record: dict[str, Any] | None = None) -> httpx.Response: + return httpx.Response( + 201, + json=record or sandbox_record(), + headers={"content-type": "application/json"}, + ) + + +def revoked_token_response(token: str) -> httpx.Response: + return httpx.Response( + 401, + json={"code": "TOKEN_REVOKED", "message": "revoked"}, + headers={ + "content-type": "application/json", + "x-blaxel-error-code": "TOKEN_REVOKED", + }, + ) + + +def test_create_mints_token_and_returns_sandbox() -> None: + transport = RoutingTransport( + routes={ + "POST /v1/sandboxes/instances": created_response(), + } + ) + client = make_sync_client(transport) + + sandbox = client.create(name="sb-1") + + assert sandbox.name == "sb-1" + assert sandbox.url == SANDBOX_URL + assert isinstance(sandbox.info, SandboxInfo) + assert sandbox.info.status == "DEPLOYED" + body = json.loads(transport.requests[1][3]) + assert body == {"name": "sb-1"} + assert transport.authorizations("/v1/sandboxes/instances") == ["Bearer tok-1"] + client.close() + + +def test_token_is_minted_once_and_cached() -> None: + transport = RoutingTransport( + routes={ + "POST /v1/sandboxes/instances": [created_response(), created_response()], + "GET /v1/sandboxes/instances/sb-1": sandbox_record(), + } + ) + client = make_sync_client(transport) + + client.create() + client.get_info("sb-1") + + assert transport.mint_count == 1 + client.close() + + +def test_empty_api_key_sends_no_authorization() -> None: + transport = RoutingTransport( + routes={"GET /v1/sandboxes/instances/sb-1": sandbox_record()} + ) + client = SandboxClient( + api_key="", + base_url_override=MANAGEMENT_URL, + transport_override=transport.sync_transport, + ) + + client.get_info("sb-1") + + assert transport.authorizations("/v1/sandboxes/instances/sb-1") == [] + client.close() + + +def test_revoked_token_is_replaced_and_request_resent() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": lambda request: ( + revoked_token_response("tok-1") + if request.headers.get("authorization") == "Bearer tok-1" + else httpx.Response( + 200, + json=sandbox_record(), + headers={"content-type": "application/json"}, + ) + ), + }, + mint_tokens=["tok-1", "tok-2"], + ) + client = make_sync_client(transport) + + info = client.get_info("sb-1") + + assert info.name == "sb-1" + assert transport.mint_count == 2 + assert transport.authorizations("/v1/sandboxes/instances/sb-1") == [ + "Bearer tok-1", + "Bearer tok-2", + ] + client.close() + + +def test_list_follows_cursor_pages() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances": [ + { + "items": [sandbox_record("sb-1")], + "pagination": {"has_more": True, "cursor": "c1"}, + }, + { + "items": [sandbox_record("sb-2")], + "pagination": {"has_more": False}, + }, + ], + } + ) + client = make_sync_client(transport) + + names = [info.name for info in client.list()] + + assert names == ["sb-1", "sb-2"] + assert transport.count("GET", "/v1/sandboxes/instances") == 2 + client.close() + + +def test_list_repeated_cursor_raises() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances": { + "items": [sandbox_record()], + "pagination": {"has_more": True, "cursor": "c1"}, + }, + } + ) + client = make_sync_client(transport) + + with pytest.raises(ValueError, match="repeated cursor"): + list(client.list()) + + client.close() + + +def test_delete_returns_record() -> None: + transport = RoutingTransport( + routes={ + "DELETE /v1/sandboxes/instances/sb-1": httpx.Response( + 202, + json=sandbox_record(status="DELETING"), + headers={"content-type": "application/json"}, + ) + } + ) + client = make_sync_client(transport) + + info = client.delete("sb-1") + + assert info.status == "DELETING" + client.close() + + +def test_control_plane_error_becomes_sandbox_api_error() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": httpx.Response( + 404, json={"code": "NOT_FOUND", "message": "no such sandbox"} + ) + } + ) + client = make_sync_client(transport) + + with pytest.raises(SandboxApiError) as exc_info: + client.get_info("sb-1") + + assert exc_info.value.status == 404 + assert exc_info.value.code == "NOT_FOUND" + client.close() + + +def make_full_client(transport: RoutingTransport) -> SandboxClient: + transport.routes.setdefault("POST /v1/sandboxes/instances", created_response()) + transport.routes.setdefault("GET /filesystem/file.txt", file_record()) + transport.routes.setdefault("PUT /filesystem/file.txt", {"message": "written"}) + transport.routes.setdefault("POST /process", process_record()) + return make_sync_client(transport) + + +def test_fs_read_and_write_on_created_sandbox() -> None: + transport = RoutingTransport(routes={}) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + assert sandbox.fs.read("file.txt") == "file text" + sandbox.fs.write("file.txt", "new text") + + assert transport.count("GET", "/filesystem/file.txt") == 1 + written = json.loads(transport.requests[-1][3]) + assert written == {"content": "new text"} + client.close() + + +def test_fs_read_retries_dropped_connection() -> None: + transport = RoutingTransport( + routes={ + "GET /filesystem/file.txt": [httpx.ConnectError("dropped"), file_record()] + } + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + content = sandbox.fs.read("file.txt") + + assert content == "file text" + assert transport.count("GET", "/filesystem/file.txt") == 2 + client.close() + + +def test_fs_read_directory_raises() -> None: + transport = RoutingTransport( + routes={ + "GET /filesystem/dir": { + "files": [], + "name": "dir", + "path": "/dir", + "subdirectories": [], + } + } + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + with pytest.raises(ValueError, match="is a directory"): + sandbox.fs.read("dir") + + client.close() + + +def test_process_exec_returns_curated_info() -> None: + transport = RoutingTransport(routes={}) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + info = sandbox.process.exec("echo hi", working_dir="/tmp", wait_for_completion=True) + + assert info.pid == "123" + assert info.status == "completed" + assert info.exit_code == 0 + assert info.working_dir == "/tmp" + assert info.completed_at is not None + assert info.completed_at.year == 2026 + body = json.loads(transport.requests[-1][3]) + assert body["command"] == "echo hi" + assert body["workingDir"] == "/tmp" + assert body["waitForCompletion"] is True + client.close() + + +def test_exec_gateway_error_is_sandbox_gateway_error() -> None: + transport = RoutingTransport( + routes={"POST /process": httpx.Response(502, text="bad gateway")} + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + with pytest.raises(SandboxGatewayError) as exc_info: + sandbox.process.exec("echo hi") + + assert exc_info.value.status == 502 + client.close() + + +def test_get_by_name_and_by_info() -> None: + transport = RoutingTransport( + routes={"GET /v1/sandboxes/instances/sb-1": sandbox_record()} + ) + client = make_sync_client(transport) + + info = client.get_info("sb-1") + by_name = client.get("sb-1") + assert by_name.name == "sb-1" + assert by_name.info is None + + by_record = client.get(info) + assert by_record.name == "sb-1" + + # The by-record get made no call of its own. + assert transport.count("GET", "/v1/sandboxes/instances/sb-1") == 2 + client.close() + + +@pytest.mark.asyncio +async def test_async_client_create_and_exec() -> None: + transport = RoutingTransport(routes={}) + transport.routes.setdefault("POST /v1/sandboxes/instances", created_response()) + transport.routes.setdefault("POST /process", process_record()) + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + + sandbox = await client.create(name="sb-1") + info = await sandbox.process.exec("echo hi") + + assert info.pid == "123" + assert transport.authorizations("/process") == ["Bearer tok-1"] + await client.close() + + +@pytest.mark.asyncio +async def test_async_list_paginates() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances": [ + { + "items": [sandbox_record("sb-1")], + "pagination": {"has_more": True, "cursor": "c1"}, + }, + { + "items": [sandbox_record("sb-2")], + "pagination": {"has_more": False}, + }, + ], + } + ) + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + + names = [info.name async for info in client.list()] + + assert names == ["sb-1", "sb-2"] + await client.close() From cab62abe09e42f39d259a010f91fd794840174c6 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Tue, 29 Sep 2026 18:04:31 -0700 Subject: [PATCH 02/18] Trim comment bloat in the sandbox surface Co-Authored-By: Claude Code --- baseten/sandbox/_auth.py | 15 ++++++--------- baseten/sandbox/_client.py | 15 +++++++-------- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py index ffe907d..3391015 100644 --- a/baseten/sandbox/_auth.py +++ b/baseten/sandbox/_auth.py @@ -8,19 +8,16 @@ import httpx -# Refresh this long before the token's stated expiry, to allow for clock skew -# between this machine and the server, and for time spent in flight. +# Refresh before the stated expiry: clock skew, plus time in flight. _TOKEN_EXPIRY_LEEWAY = timedelta(seconds=60) -# A few retries cover a new token being invalidated in the same event as the -# old one. Past that, the failure is raised, since tokens that keep getting -# rejected point at something a new token cannot fix. +# A few retries cover a new token invalidated in the same event as the old +# one; beyond that, rejection points at something a new token cannot fix. _TOKEN_INVALIDATION_MAX_RETRIES = 2 -# The control plane and sandboxes send this header when they reject a request's -# token as revoked, which can happen before the token's stated expiry. The -# request was rejected before any work was done, so sending it again with a -# new token is safe even when it has side effects. +# Both planes send this header on revocation, which can precede the stated +# expiry. The request was untouched, so re-sending with a new token is safe +# even when it has side effects. _TOKEN_REVOKED_CODE = "TOKEN_REVOKED" diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py index b0b18ae..c9fd268 100644 --- a/baseten/sandbox/_client.py +++ b/baseten/sandbox/_client.py @@ -30,8 +30,8 @@ DEFAULT_MANAGEMENT_BASE_URL = "https://api.baseten.co" -# Only connect time is bounded by default: a sandbox call can run for as long -# as the operation needs, and callers pass a stricter `timeout` when wanted. +# Only connect time is bounded: a sandbox call runs as long as the operation +# needs. Callers pass a stricter `timeout` when wanted. DEFAULT_TIMEOUT = httpx.Timeout(None, connect=10.0) @@ -166,16 +166,15 @@ def __init__( http2=resolve_http2_flag(self._options.http2) ) elif isinstance(transport_override, httpx.BaseTransport): - # MockTransport implements both interfaces, so an async-only - # transport (not a BaseTransport) is what must fall through. + # MockTransport implements both interfaces; only an async-only + # transport fails this isinstance. self._pool = transport_override else: raise TypeError( "SandboxClient needs a sync transport in transport_override" ) self._timeout = self._options.timeout or DEFAULT_TIMEOUT - # The mint client authenticates with the API key directly; every - # other client goes through the auth transport below. + # API-key auth, not the auth transport: minting must not need a token. self._mint_http_client = httpx.Client( transport=self._pool, base_url=_management_base_url(self._options), @@ -200,8 +199,8 @@ def __init__( self._sandbox_http_clients: list[httpx.Client] = [] def _mint(self) -> tuple[str, datetime]: - # The mint client authenticates with the API key, not with a token, - # so this does not go back through the auth transport. + # The mint client carries the API key, so minting cannot recurse + # through the auth transport. minted = baseten.client.managementapi.ApiClient( self._mint_http_client ).post_token( From 2c82c851ae60714a73150b79003f6041b52e8c2d Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Tue, 29 Sep 2026 18:23:27 -0700 Subject: [PATCH 03/18] Pin curated-surface translation with parity tests The surface exposes no generated types: methods take keyword arguments of primitives and curated dataclasses, and the generated models are built internally. The parity tests fail when an upstream spec field appears, until the surface forwards it or lists it as deliberately not translated. Co-Authored-By: Claude Code --- tests/sandbox/test_parity.py | 118 +++++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 tests/sandbox/test_parity.py diff --git a/tests/sandbox/test_parity.py b/tests/sandbox/test_parity.py new file mode 100644 index 0000000..03b592f --- /dev/null +++ b/tests/sandbox/test_parity.py @@ -0,0 +1,118 @@ +"""Parity between the curated surface and the generated models it wraps. + +The surface never exposes generated types: methods take keyword arguments of +primitives and curated dataclasses, and the generated models are constructed +internally. These tests pin that translation, so a spec field added upstream +fails here until the surface either forwards it or lists it under +INTENTIONALLY_NOT_TRANSLATED_*. +""" + +from __future__ import annotations + +import baseten.client.managementapi +import baseten.client.sandboxapi + +# Control-plane fields left to the raw client. +INTENTIONALLY_NOT_TRANSLATED_SANDBOX = {"lifecycle", "network", "ports"} + +INTENTIONALLY_NOT_TRANSLATED_PROCESS = { + "keepAlive", + "maxRestarts", + "restartOnFailure", + "stdin", + "waitForPorts", +} + +# Deliberately matching the JS surface's process record. +INTENTIONALLY_NOT_TRANSLATED_PROCESS_RECORD = { + "keepAlive", + "maxRestarts", + "restartCount", + "restartOnFailure", + "stdin", +} + +# File-write fields left to the raw client. +INTENTIONALLY_NOT_TRANSLATED_FILE = {"isDirectory", "permissions"} + + +def test_create_forwards_every_translated_field() -> None: + forwarded = { + "name", + "image", + "memory", + "region", + "envs", + "labels", + "display_name", + "external_id", + } + fields = set(baseten.client.managementapi.CreateSandboxRequest.model_fields) + assert forwarded <= fields + assert fields - forwarded == INTENTIONALLY_NOT_TRANSLATED_SANDBOX + + +def test_record_translation_covers_every_translated_field() -> None: + translated = { + "name", + "url", + "status", + "state", + "image", + "memory", + "region", + "enabled", + "envs", + "labels", + "display_name", + "external_id", + "created_at", + "updated_at", + "created_by", + "updated_by", + "last_used_at", + "expires_in", + } + fields = set(baseten.client.managementapi.Sandbox.model_fields) + assert translated <= fields + assert fields - translated == INTENTIONALLY_NOT_TRANSLATED_SANDBOX + + +def test_exec_forwards_every_translated_field() -> None: + forwarded = { + "command", + "workingDir", + "env", + "name", + "timeout", + "waitForCompletion", + } + fields = set(baseten.client.sandboxapi.ProcessRequest.model_fields) + assert forwarded <= fields + assert fields - forwarded == INTENTIONALLY_NOT_TRANSLATED_PROCESS + + +def test_process_record_translation_covers_every_translated_field() -> None: + translated = { + "command", + "completedAt", + "exitCode", + "logs", + "name", + "pid", + "startedAt", + "status", + "stderr", + "stdout", + "workingDir", + } + fields = set(baseten.client.sandboxapi.ProcessResponse.model_fields) + assert translated <= fields + assert fields - translated == INTENTIONALLY_NOT_TRANSLATED_PROCESS_RECORD + + +def test_fs_write_forwards_every_translated_field() -> None: + forwarded = {"content"} + fields = set(baseten.client.sandboxapi.FileRequest.model_fields) + assert forwarded <= fields + assert fields - forwarded == INTENTIONALLY_NOT_TRANSLATED_FILE From f305648dd34ad54fa7465916821f377ed361910c Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Tue, 29 Sep 2026 18:47:01 -0700 Subject: [PATCH 04/18] Add sandboxes_base_url_override and a live e2e suite The sandbox management API can be served on a different host than the token exchange (today: api.blaxel.dev serves /v1/sandboxes while api.staging.baseten.co mints the tokens, until staging.b10.co is live), so the client accepts a separate base URL for the sandboxes calls. The e2e suite mirrors the JS suite's environment contract (BASETEN_E2E_TEST_API_KEY, BASETEN_E2E_TEST_DOMAIN, BASETEN_E2E_TEST_SANDBOXES_DOMAIN) and skips without them. Verified live against staging + blaxel dev: create (server applied its image and memory defaults), exec with waitForCompletion, fs write/read, delete. Co-Authored-By: Claude Code --- baseten/sandbox/_client.py | 50 +++++++++++++--- tests/sandbox/test_sandbox_e2e.py | 94 +++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 8 deletions(-) create mode 100644 tests/sandbox/test_sandbox_e2e.py diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py index c9fd268..33efca8 100644 --- a/baseten/sandbox/_client.py +++ b/baseten/sandbox/_client.py @@ -58,6 +58,10 @@ class SandboxClientOptions: base_url_override: str | None = None """Explicit management API base URL, or None to use the default.""" + sandboxes_base_url_override: str | None = None + """Where /v1/sandboxes is served, when that differs from the management + API. The token exchange still goes to the management API.""" + headers: Mapping[str, str] | None = None """Additional headers to send on every request.""" @@ -139,6 +143,7 @@ def __init__( token_provider: SandboxTokenProvider | None = None, team_id: str | None = None, base_url_override: str | None = None, + sandboxes_base_url_override: str | None = None, headers: Mapping[str, str] | None = None, http2: bool | None = None, timeout: httpx.Timeout | None = None, @@ -155,6 +160,7 @@ def __init__( token_provider=token_provider, team_id=team_id, base_url_override=base_url_override, + sandboxes_base_url_override=sandboxes_base_url_override, headers=headers, http2=http2, timeout=timeout, @@ -196,6 +202,17 @@ def __init__( timeout=self._timeout, ) self._api = baseten.client.managementapi.ApiClient(self._http_client) + self._sandbox_api = self._api + if self._options.sandboxes_base_url_override is not None: + self._sandboxes_http_client = httpx.Client( + transport=self._auth_transport, + base_url=self._options.sandboxes_base_url_override, + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._sandbox_api = baseten.client.managementapi.ApiClient( + self._sandboxes_http_client + ) self._sandbox_http_clients: list[httpx.Client] = [] def _mint(self) -> tuple[str, datetime]: @@ -242,7 +259,7 @@ def create( for what the caller leaves unset. """ try: - sandbox = self._api.create_sandbox( + sandbox = self._sandbox_api.create_sandbox( params=baseten.client.managementapi.CreateSandboxParams( team_id=self._options.team_id ), @@ -265,7 +282,7 @@ def create( def get_info(self, name: str) -> SandboxInfo: """Get a sandbox's current record.""" try: - sandbox = self._api.get_sandbox( + sandbox = self._sandbox_api.get_sandbox( sandbox_name=name, params=baseten.client.managementapi.GetSandboxParams( team_id=self._options.team_id @@ -297,7 +314,7 @@ def list( cursor: str | None = None while True: try: - response = self._api.list_sandboxes( + response = self._sandbox_api.list_sandboxes( params=baseten.client.managementapi.ListSandboxesParams( team_id=self._options.team_id, cursor=cursor, @@ -331,7 +348,7 @@ def delete(self, name: str) -> SandboxInfo: shows the sandbox as still deleting. """ try: - sandbox = self._api.delete_sandbox( + sandbox = self._sandbox_api.delete_sandbox( sandbox_name=name, params=baseten.client.managementapi.DeleteSandboxParams( team_id=self._options.team_id @@ -345,6 +362,8 @@ def close(self) -> None: """Close the client and the connections every sandbox shares.""" for http_client in self._sandbox_http_clients: http_client.close() + if self._options.sandboxes_base_url_override is not None: + self._sandboxes_http_client.close() self._http_client.close() self._mint_http_client.close() @@ -391,6 +410,7 @@ def __init__( token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None = None, team_id: str | None = None, base_url_override: str | None = None, + sandboxes_base_url_override: str | None = None, headers: Mapping[str, str] | None = None, http2: bool | None = None, timeout: httpx.Timeout | None = None, @@ -408,6 +428,7 @@ def __init__( token_provider=token_provider, team_id=team_id, base_url_override=base_url_override, + sandboxes_base_url_override=sandboxes_base_url_override, headers=headers, http2=http2, timeout=timeout, @@ -457,6 +478,17 @@ async def mint() -> tuple[str, datetime]: timeout=self._timeout, ) self._api = baseten.client.managementapi.AsyncApiClient(self._http_client) + self._sandbox_api = self._api + if self._options.sandboxes_base_url_override is not None: + self._sandboxes_http_client = httpx.AsyncClient( + transport=self._auth_transport, + base_url=self._options.sandboxes_base_url_override, + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._sandbox_api = baseten.client.managementapi.AsyncApiClient( + self._sandboxes_http_client + ) self._sandbox_http_clients: list[httpx.AsyncClient] = [] @property @@ -491,7 +523,7 @@ async def create( for what the caller leaves unset. """ try: - sandbox = await self._api.create_sandbox( + sandbox = await self._sandbox_api.create_sandbox( params=baseten.client.managementapi.CreateSandboxParams( team_id=self._options.team_id ), @@ -514,7 +546,7 @@ async def create( async def get_info(self, name: str) -> SandboxInfo: """Get a sandbox's current record.""" try: - sandbox = await self._api.get_sandbox( + sandbox = await self._sandbox_api.get_sandbox( sandbox_name=name, params=baseten.client.managementapi.GetSandboxParams( team_id=self._options.team_id @@ -546,7 +578,7 @@ async def list( cursor: str | None = None while True: try: - response = await self._api.list_sandboxes( + response = await self._sandbox_api.list_sandboxes( params=baseten.client.managementapi.ListSandboxesParams( team_id=self._options.team_id, cursor=cursor, @@ -581,7 +613,7 @@ async def delete(self, name: str) -> SandboxInfo: shows the sandbox as still deleting. """ try: - sandbox = await self._api.delete_sandbox( + sandbox = await self._sandbox_api.delete_sandbox( sandbox_name=name, params=baseten.client.managementapi.DeleteSandboxParams( team_id=self._options.team_id @@ -595,6 +627,8 @@ async def close(self) -> None: """Close the client and the connections every sandbox shares.""" for http_client in self._sandbox_http_clients: await http_client.aclose() + if self._options.sandboxes_base_url_override is not None: + await self._sandboxes_http_client.aclose() await self._http_client.aclose() await self._mint_http_client.aclose() diff --git a/tests/sandbox/test_sandbox_e2e.py b/tests/sandbox/test_sandbox_e2e.py new file mode 100644 index 0000000..f8a1313 --- /dev/null +++ b/tests/sandbox/test_sandbox_e2e.py @@ -0,0 +1,94 @@ +"""End-to-end tests against a live environment. + +Skipped unless BASETEN_E2E_TEST_API_KEY and BASETEN_E2E_TEST_DOMAIN are set. +The key must belong to an organization with sandboxes enabled. +BASETEN_E2E_TEST_SANDBOXES_DOMAIN optionally points the /v1/sandboxes calls +at a different host than the token exchange, which stays on the management +domain. Mirrors the JS suite's environment contract. +""" + +from __future__ import annotations + +import os +import time + +import pytest + +from baseten.sandbox import AsyncSandboxClient, SandboxClient + +API_KEY = os.environ.get("BASETEN_E2E_TEST_API_KEY", "") +MANAGEMENT_DOMAIN = os.environ.get("BASETEN_E2E_TEST_DOMAIN", "") +SANDBOXES_DOMAIN = ( + os.environ.get("BASETEN_E2E_TEST_SANDBOXES_DOMAIN") or MANAGEMENT_DOMAIN +) + +e2e_required = pytest.mark.skipif( + not (API_KEY and MANAGEMENT_DOMAIN), reason="e2e environment not set" +) + +# The one region where the execution plane accepts staging-minted tokens +# (eu-dub-1 rejects them with AUTHENTICATION_FAILED). +E2E_REGION = "us-was-1" + + +def unique_sandbox_name() -> str: + return f"baseten-python-e2e-{time.time():.0f}" + + +def make_client() -> SandboxClient: + return SandboxClient( + api_key=API_KEY, + base_url_override=f"https://{MANAGEMENT_DOMAIN}", + sandboxes_base_url_override=( + f"https://{SANDBOXES_DOMAIN}" + if SANDBOXES_DOMAIN != MANAGEMENT_DOMAIN + else None + ), + ) + + +@e2e_required +def test_create_exec_fs_delete() -> None: + name = unique_sandbox_name() + client = make_client() + try: + sandbox = client.create(name=name, region=E2E_REGION) + assert sandbox.info is not None + assert sandbox.info.status == "DEPLOYED" + assert sandbox.info.url is not None + + process = sandbox.process.exec("echo hello", wait_for_completion=True) + assert process.status == "completed" + assert process.exit_code == 0 + assert process.stdout == "hello\n" + + sandbox.fs.write("/tmp/e2e.txt", "python sdk") + assert sandbox.fs.read("/tmp/e2e.txt") == "python sdk" + finally: + deleted = client.delete(name) + assert deleted.status in ("DELETING", "TERMINATED") + client.close() + + +@e2e_required +@pytest.mark.asyncio +async def test_async_create_and_list() -> None: + name = unique_sandbox_name() + client = AsyncSandboxClient( + api_key=API_KEY, + base_url_override=f"https://{MANAGEMENT_DOMAIN}", + sandboxes_base_url_override=( + f"https://{SANDBOXES_DOMAIN}" + if SANDBOXES_DOMAIN != MANAGEMENT_DOMAIN + else None + ), + ) + try: + sandbox = await client.create(name=name, region=E2E_REGION) + assert sandbox.name == name + + listed = [info.name async for info in client.list(query=name)] + assert name in listed + finally: + await client.delete(name) + await client.close() From 4d587f57c5315b612f5bcc966c5a43d5568b4a9f Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Tue, 29 Sep 2026 19:24:51 -0700 Subject: [PATCH 05/18] Extend the live e2e suite: get, pagination, logs stream, multipart, archive Exercises the raw execution client against a live sandbox: stdin piping with delete_process_stdin and the text/plain logs stream, multipart initiate/part/complete followed by an fs read-back, and a dry-run archive export (the 200 ExportResult branch; the 202 branch needs a presigned upload URL, which stays unit-tested). The stdin test is xfail: the server sends logs: null on a still-running process while the spec declares logs required, so pydantic rejects the response. Reported for a server-side fix. Co-Authored-By: Claude Code --- tests/sandbox/test_sandbox_e2e.py | 143 ++++++++++++++++++++++++++++++ 1 file changed, 143 insertions(+) diff --git a/tests/sandbox/test_sandbox_e2e.py b/tests/sandbox/test_sandbox_e2e.py index f8a1313..bc7564e 100644 --- a/tests/sandbox/test_sandbox_e2e.py +++ b/tests/sandbox/test_sandbox_e2e.py @@ -14,6 +14,7 @@ import pytest +import baseten.client.sandboxapi from baseten.sandbox import AsyncSandboxClient, SandboxClient API_KEY = os.environ.get("BASETEN_E2E_TEST_API_KEY", "") @@ -92,3 +93,145 @@ async def test_async_create_and_list() -> None: finally: await client.delete(name) await client.close() + + +@e2e_required +def test_get_info_and_get() -> None: + name = unique_sandbox_name() + client = make_client() + try: + created = client.create(name=name, region=E2E_REGION) + info = client.get_info(name) + assert info.name == name + assert info.status == "DEPLOYED" + # The record keeps the server's trailing slash; Sandbox.url normalizes it. + assert info.url is not None + assert info.url.rstrip("/") == created.url + + sandbox = client.get(name) + assert sandbox.name == name + assert sandbox.url == created.url + finally: + client.delete(name) + client.close() + + +@e2e_required +def test_list_paginates_across_pages() -> None: + prefix = f"baseten-python-e2e-{time.time():.0f}" + names = [f"{prefix}-a", f"{prefix}-b"] + client = make_client() + try: + for sandbox_name in names: + client.create(name=sandbox_name, region=E2E_REGION) + + listed = [info.name for info in client.list(query=prefix, page_size=1)] + assert sorted(listed) == sorted(names) + finally: + for sandbox_name in names: + client.delete(sandbox_name) + client.close() + + +@e2e_required +@pytest.mark.xfail( + reason="server bug: POST /process with waitForCompletion=false sends" + " logs: null, but the spec declares logs required; pydantic rejects the" + " response (JS does not validate, so only Python fails). Tracked with Chad.", + strict=True, +) +def test_logs_stream_and_stdin_via_raw_api() -> None: + name = unique_sandbox_name() + client = make_client() + try: + sandbox = client.create(name=name, region=E2E_REGION) + # `cat` with a stdin pipe, driven through the raw execution client: + # the surface's exec does not expose stdin yet. + sandbox.api.post_process( + request=baseten.client.sandboxapi.ProcessRequest( + command="cat", stdin=True, name="piper", waitForCompletion=False + ) + ) + sandbox.api.post_process_stdin( + identifier="piper", content=b"stdin round trip\n" + ) + # Closing stdin makes cat exit; closing again is the no-op probe. + sandbox.api.delete_process_stdin(identifier="piper") + for _ in range(20): + time.sleep(0.5) + process = sandbox.api.get_process_identifier(identifier="piper") + if process.status == "completed": + assert "stdin round trip" in process.stdout + break + else: + raise AssertionError("cat did not complete after stdin closed") + + logs_response = sandbox.api.get_process_logs_stream(identifier="piper") + logs_response.read() + logs_response.close() + assert "stdin round trip" in logs_response.text + finally: + client.delete(name) + client.close() + + +@e2e_required +def test_multipart_upload_via_raw_api() -> None: + name = unique_sandbox_name() + client = make_client() + try: + sandbox = client.create(name=name, region=E2E_REGION) + remote_path = "/tmp/multipart-e2e.txt" + initiated = sandbox.api.post_filesystem_multipart_initiate( + path=remote_path, + request=baseten.client.sandboxapi.MultipartInitiateRequest(), + ) + upload_id = initiated.uploadId + assert upload_id is not None + + part = sandbox.api.put_filesystem_multipart_part( + upload_id=upload_id, + params=baseten.client.sandboxapi.PutFilesystemMultipartPartParams( + partNumber=1 + ), + files={"file": ("part.bin", b"multipart e2e content")}, + ) + assert part.partNumber == 1 + + completed = sandbox.api.post_filesystem_multipart_complete( + upload_id=upload_id, + request=baseten.client.sandboxapi.MultipartCompleteRequest( + parts=[ + baseten.client.sandboxapi.MultipartPartInfo( + partNumber=part.partNumber, etag=part.etag + ) + ] + ), + ) + assert completed.message + + assert sandbox.fs.read(remote_path) == "multipart e2e content" + finally: + client.delete(name) + client.close() + + +@e2e_required +def test_archive_export_status_dispatch() -> None: + name = unique_sandbox_name() + client = make_client() + try: + sandbox = client.create(name=name, region=E2E_REGION) + sandbox.fs.write("/tmp/archive-e2e.txt", "archive me") + + result = sandbox.api.post_archive_export( + request=baseten.client.sandboxapi.ExportOptions(dryRun=True) + ) + assert isinstance(result, baseten.client.sandboxapi.ExportResult) + assert result.manifest is not None + + # The async (202 ExportProgress) path needs a presigned upload URL, + # which the test cannot mint; it stays covered by unit tests. + finally: + client.delete(name) + client.close() From 6a3b2aad8b981e768211377bb81954ce1f085197 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Wed, 30 Sep 2026 12:27:41 -0700 Subject: [PATCH 06/18] Gate the sandbox e2e suite on the sandboxes domain being set CI runs the suite off the repo's existing e2e secrets, but no environment serves /v1/sandboxes on the token-exchange domain, so without the explicit split host the suite fails on DNS or 404s. The suite now runs only when BASETEN_E2E_TEST_SANDBOXES_DOMAIN is set alongside the other two secrets, and skips otherwise. Co-Authored-By: Claude Code --- tests/sandbox/test_sandbox_e2e.py | 30 +++++++++++------------------- 1 file changed, 11 insertions(+), 19 deletions(-) diff --git a/tests/sandbox/test_sandbox_e2e.py b/tests/sandbox/test_sandbox_e2e.py index bc7564e..aa85006 100644 --- a/tests/sandbox/test_sandbox_e2e.py +++ b/tests/sandbox/test_sandbox_e2e.py @@ -1,10 +1,11 @@ """End-to-end tests against a live environment. -Skipped unless BASETEN_E2E_TEST_API_KEY and BASETEN_E2E_TEST_DOMAIN are set. -The key must belong to an organization with sandboxes enabled. -BASETEN_E2E_TEST_SANDBOXES_DOMAIN optionally points the /v1/sandboxes calls -at a different host than the token exchange, which stays on the management -domain. Mirrors the JS suite's environment contract. +Skipped unless all of BASETEN_E2E_TEST_API_KEY, BASETEN_E2E_TEST_DOMAIN, +and BASETEN_E2E_TEST_SANDBOXES_DOMAIN are set. The key must belong to an +organization with sandboxes enabled, and the sandbox management API must +be served on its own host: today no environment serves /v1/sandboxes on +the token-exchange domain, so the split is required rather than optional. +Mirrors the JS suite's environment contract. """ from __future__ import annotations @@ -19,12 +20,11 @@ API_KEY = os.environ.get("BASETEN_E2E_TEST_API_KEY", "") MANAGEMENT_DOMAIN = os.environ.get("BASETEN_E2E_TEST_DOMAIN", "") -SANDBOXES_DOMAIN = ( - os.environ.get("BASETEN_E2E_TEST_SANDBOXES_DOMAIN") or MANAGEMENT_DOMAIN -) +SANDBOXES_DOMAIN = os.environ.get("BASETEN_E2E_TEST_SANDBOXES_DOMAIN", "") e2e_required = pytest.mark.skipif( - not (API_KEY and MANAGEMENT_DOMAIN), reason="e2e environment not set" + not (API_KEY and MANAGEMENT_DOMAIN and SANDBOXES_DOMAIN), + reason="sandbox e2e environment not set", ) # The one region where the execution plane accepts staging-minted tokens @@ -40,11 +40,7 @@ def make_client() -> SandboxClient: return SandboxClient( api_key=API_KEY, base_url_override=f"https://{MANAGEMENT_DOMAIN}", - sandboxes_base_url_override=( - f"https://{SANDBOXES_DOMAIN}" - if SANDBOXES_DOMAIN != MANAGEMENT_DOMAIN - else None - ), + sandboxes_base_url_override=f"https://{SANDBOXES_DOMAIN}", ) @@ -78,11 +74,7 @@ async def test_async_create_and_list() -> None: client = AsyncSandboxClient( api_key=API_KEY, base_url_override=f"https://{MANAGEMENT_DOMAIN}", - sandboxes_base_url_override=( - f"https://{SANDBOXES_DOMAIN}" - if SANDBOXES_DOMAIN != MANAGEMENT_DOMAIN - else None - ), + sandboxes_base_url_override=f"https://{SANDBOXES_DOMAIN}", ) try: sandbox = await client.create(name=name, region=E2E_REGION) From 4b34014684d171084787a0d560117d74843639c6 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Wed, 30 Sep 2026 12:51:04 -0700 Subject: [PATCH 07/18] Fix review findings: mint atomicity, revocation exhaustion, gateway retries - The token cache locks the miss check and the mint as one sequence, so concurrent callers on a cold cache send a single exchange. - Auth-transport revocation retries return the last 401 response when exhausted, instead of raising AssertionError; the generated client now raises a meaningful error for a persistently revoked token. - The auth transports delegate close() to the inner pool, so standalone sandboxes and their owners no longer leak the connection pool. - fs.read converts each attempt's error inside the retried callable, so gateway 502/503/504 responses actually retry; previously conversion happened after the retry helper had already given up. - create() omits name when unset instead of sending null; the mint client honors the configured timeout; typed error models dump to dicts so the curated error reads code and details; raw sandbox access is exposed as .sandbox_api when the domain split moves /v1/sandboxes off .api; the standalone async sandbox accepts async token providers in its annotation. Co-Authored-By: Claude Code --- baseten/sandbox/_auth.py | 70 +++++++++++++++++++-------------- baseten/sandbox/_client.py | 28 +++++++++++++- baseten/sandbox/_errors.py | 5 +++ baseten/sandbox/_sandbox.py | 47 ++++++++++++++-------- tests/sandbox/test_surface.py | 73 +++++++++++++++++++++++++++++++++++ 5 files changed, 175 insertions(+), 48 deletions(-) diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py index 3391015..1403db7 100644 --- a/baseten/sandbox/_auth.py +++ b/baseten/sandbox/_auth.py @@ -54,28 +54,30 @@ def resolve_http2_flag(http2: bool | None) -> bool: class _TokenCache: - """A cached minted token, guarded so concurrent callers share one mint.""" + """A cached minted token, guarded so concurrent callers share one mint. + + The lock spans the miss check and the mint, so a burst on a cold cache + sends a single exchange. + """ def __init__(self) -> None: self._lock = threading.Lock() self._token: str | None = None self._expires_at: datetime | None = None - def get(self) -> tuple[str, datetime] | None: + def get_or_mint( + self, mint: Callable[[], tuple[str, datetime]] + ) -> tuple[str, datetime]: with self._lock: - if self._token is None or self._expires_at is None: - return None - if datetime.now(UTC) >= self._expires_at - _TOKEN_EXPIRY_LEEWAY: - self._token = None - self._expires_at = None - return None + if ( + self._token is not None + and self._expires_at is not None + and datetime.now(UTC) < self._expires_at - _TOKEN_EXPIRY_LEEWAY + ): + return self._token, self._expires_at + self._token, self._expires_at = mint() return self._token, self._expires_at - def put(self, token: str, expires_at: datetime) -> None: - with self._lock: - self._token = token - self._expires_at = expires_at - def drop(self, token: str) -> None: """Drop a token the server rejected, so the next request gets a new one. @@ -118,11 +120,7 @@ def token(self) -> str | None: return self._provider() if self._mint is None: return None - cached = self._cache.get() - if cached is not None: - return cached[0] - token, expires_at = self._mint() - self._cache.put(token, expires_at) + token, _ = self._cache.get_or_mint(self._mint) return token def invalidate(self, token: str) -> None: @@ -194,21 +192,27 @@ def __init__( def handle_request(self, request: httpx.Request) -> httpx.Response: request.read() - for _ in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): + for attempt in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): token = self._token_source.token() if token is not None: request.headers["Authorization"] = f"Bearer {token}" response = self._inner.handle_request(request) - if ( - token is None - or response.status_code != 401 - or response.headers.get("x-blaxel-error-code") != _TOKEN_REVOKED_CODE - ): + revoked = ( + token is not None + and response.status_code == 401 + and response.headers.get("x-blaxel-error-code") == _TOKEN_REVOKED_CODE + ) + # The last attempt's response is returned even when still + # revoked, so the generated client raises a meaningful error. + if not revoked or attempt == _TOKEN_INVALIDATION_MAX_RETRIES: return response self._token_source.invalidate(token) response.close() raise AssertionError # pragma: no cover - loop returns first + def close(self) -> None: + self._inner.close() + class _AsyncAuthTransport(httpx.AsyncBaseTransport): """Async variant of :class:`_SyncAuthTransport`.""" @@ -221,17 +225,23 @@ def __init__( async def handle_async_request(self, request: httpx.Request) -> httpx.Response: request.read() - for _ in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): + for attempt in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): token = await self._token_source.token() if token is not None: request.headers["Authorization"] = f"Bearer {token}" response = await self._inner.handle_async_request(request) - if ( - token is None - or response.status_code != 401 - or response.headers.get("x-blaxel-error-code") != _TOKEN_REVOKED_CODE - ): + revoked = ( + token is not None + and response.status_code == 401 + and response.headers.get("x-blaxel-error-code") == _TOKEN_REVOKED_CODE + ) + # The last attempt's response is returned even when still + # revoked, so the generated client raises a meaningful error. + if not revoked or attempt == _TOKEN_INVALIDATION_MAX_RETRIES: return response await self._token_source.invalidate(token) await response.aclose() raise AssertionError # pragma: no cover - loop returns first + + async def aclose(self) -> None: + await self._inner.aclose() diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py index 33efca8..e1a9781 100644 --- a/baseten/sandbox/_client.py +++ b/baseten/sandbox/_client.py @@ -110,7 +110,9 @@ def _sandbox_create_request( mean "clear" on updates. Creation must send nothing instead, so the server-side defaults apply. """ - create_fields: dict[str, Any] = {"name": name} + create_fields: dict[str, Any] = {} + if name is not None: + create_fields["name"] = name if image is not None: create_fields["image"] = image if memory is not None: @@ -188,6 +190,7 @@ def __init__( self._options, f"Bearer {api_key}" if api_key != "" else None, ), + timeout=self._timeout, ) self._token_source = SyncTokenSource( api_key=api_key, @@ -241,6 +244,17 @@ def api(self) -> baseten.client.managementapi.ApiClient: """ return self._api + @property + def sandbox_api(self) -> baseten.client.managementapi.ApiClient: + """The generated client bound to where /v1/sandboxes is served. + + With no override set this is :attr:`api`; with one set, raw sandbox + operations go to the override host while :attr:`api` stays on the + management domain. The generated API surface is not covered by + stability guarantees and may change between versions. + """ + return self._sandbox_api + def create( self, *, @@ -448,6 +462,7 @@ def __init__( self._timeout = self._options.timeout or DEFAULT_TIMEOUT self._mint_http_client = httpx.AsyncClient( transport=self._pool, + timeout=self._timeout, base_url=_management_base_url(self._options), headers=_request_headers( self._options, @@ -505,6 +520,17 @@ def api(self) -> baseten.client.managementapi.AsyncApiClient: """ return self._api + @property + def sandbox_api(self) -> baseten.client.managementapi.AsyncApiClient: + """The generated client bound to where /v1/sandboxes is served. + + With no override set this is :attr:`api`; with one set, raw sandbox + operations go to the override host while :attr:`api` stays on the + management domain. The generated API surface is not covered by + stability guarantees and may change between versions. + """ + return self._sandbox_api + async def create( self, *, diff --git a/baseten/sandbox/_errors.py b/baseten/sandbox/_errors.py index 64ac7b6..46cb8e5 100644 --- a/baseten/sandbox/_errors.py +++ b/baseten/sandbox/_errors.py @@ -2,6 +2,8 @@ import json +from pydantic import BaseModel + import baseten.client.managementapi import baseten.client.sandboxapi @@ -83,6 +85,9 @@ def to_sandbox_api_error(error: object, plane: str) -> Exception: body = next(iter(model_fields.values())) else: # pragma: no cover - generated errors carry exactly one model body = model_fields + if isinstance(body, BaseModel): + # The curated error reads code and details off a plain dict. + body = body.model_dump() if plane == "exec" and status in _GATEWAY_STATUSES: return SandboxGatewayError(status, body) return SandboxApiError(status, body) diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py index fdf26a3..454b365 100644 --- a/baseten/sandbox/_sandbox.py +++ b/baseten/sandbox/_sandbox.py @@ -8,6 +8,7 @@ import baseten.client.sandboxapi from baseten.client._user_agent import with_user_agent from baseten.sandbox._auth import ( + AsyncSandboxTokenProvider, AsyncTokenSource, SandboxTokenProvider, SyncTokenSource, @@ -206,7 +207,7 @@ def __init__( *, name: str, url: str, - token_provider: SandboxTokenProvider | None = None, + token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None = None, headers: Mapping[str, str] | None = None, http2: bool | None = None, timeout: httpx.Timeout | None = None, @@ -355,14 +356,20 @@ def read(self, path: str) -> str: Retries dropped connections and edge gateway errors, since a read is safe to repeat. """ - try: - result = retry_idempotent( - lambda: self._api.get_filesystem(path=path), - max_retries=self._retries.read_max_retries, - gateway_max_retries=self._retries.gateway_max_retries, - ) - except Exception as error: - raise to_sandbox_api_error(error, "exec") from error + + def attempt() -> baseten.client.sandboxapi.GetFilesystemResponse: + # Errors convert per attempt, so the retry helper sees + # SandboxGatewayError instead of the generated ResponseError. + try: + return self._api.get_filesystem(path=path) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + result = retry_idempotent( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) root = result.root if isinstance(root, baseten.client.sandboxapi.FileWithContent): return root.content @@ -398,14 +405,20 @@ async def read(self, path: str) -> str: Retries dropped connections and edge gateway errors, since a read is safe to repeat. """ - try: - result = await retry_idempotent_async( - lambda: self._api.get_filesystem(path=path), - max_retries=self._retries.read_max_retries, - gateway_max_retries=self._retries.gateway_max_retries, - ) - except Exception as error: - raise to_sandbox_api_error(error, "exec") from error + + async def attempt() -> baseten.client.sandboxapi.GetFilesystemResponse: + # Errors convert per attempt, so the retry helper sees + # SandboxGatewayError instead of the generated ResponseError. + try: + return await self._api.get_filesystem(path=path) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + result = await retry_idempotent_async( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) root = result.root if isinstance(root, baseten.client.sandboxapi.FileWithContent): return root.content diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py index 13ec95d..56dde33 100644 --- a/tests/sandbox/test_surface.py +++ b/tests/sandbox/test_surface.py @@ -467,3 +467,76 @@ async def test_async_list_paginates() -> None: assert names == ["sb-1", "sb-2"] await client.close() + + +def test_create_without_name_omits_the_field() -> None: + transport = RoutingTransport( + routes={"POST /v1/sandboxes/instances": created_response()} + ) + client = make_sync_client(transport) + + client.create() + + body = json.loads(transport.requests[1][3]) + assert body == {} + client.close() + + +def test_fs_read_retries_gateway_errors() -> None: + transport = RoutingTransport( + routes={ + "GET /filesystem/file.txt": [ + httpx.Response(502, text="bad gateway"), + httpx.Response(503, text="unavailable"), + file_record(), + ] + } + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + content = sandbox.fs.read("file.txt") + + assert content == "file text" + assert transport.count("GET", "/filesystem/file.txt") == 3 + client.close() + + +def test_persistent_revocation_returns_last_response() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": lambda request: revoked_token_response( + "revoked" + ) + }, + mint_tokens=["tok-1", "tok-2", "tok-3"], + ) + client = make_sync_client(transport) + + with pytest.raises(SandboxApiError) as exc_info: + client.get_info("sb-1") + + # Exhaustion surfaces the final 401 as a converted error, and every + # minted token was tried. + assert exc_info.value.status == 401 + assert transport.mint_count == 3 + client.close() + + +def test_typed_error_body_becomes_a_dict() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": httpx.Response( + 404, + json={"code": "NOT_FOUND", "message": "no such sandbox"}, + ) + } + ) + client = make_sync_client(transport) + + with pytest.raises(SandboxApiError) as exc_info: + client.get_info("sb-1") + + assert exc_info.value.code == "NOT_FOUND" + assert isinstance(exc_info.value.body, dict) + client.close() From 1c5b097cd5b01210fe7056a68f4d62c6d6bbdef7 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Wed, 30 Sep 2026 12:55:39 -0700 Subject: [PATCH 08/18] Trim a comment to its why Co-Authored-By: Claude Code --- baseten/sandbox/_sandbox.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py index 454b365..8cbab4e 100644 --- a/baseten/sandbox/_sandbox.py +++ b/baseten/sandbox/_sandbox.py @@ -358,8 +358,8 @@ def read(self, path: str) -> str: """ def attempt() -> baseten.client.sandboxapi.GetFilesystemResponse: - # Errors convert per attempt, so the retry helper sees - # SandboxGatewayError instead of the generated ResponseError. + # Convert per attempt: the retry helper must see + # SandboxGatewayError, not the generated ResponseError. try: return self._api.get_filesystem(path=path) except Exception as error: @@ -407,8 +407,8 @@ async def read(self, path: str) -> str: """ async def attempt() -> baseten.client.sandboxapi.GetFilesystemResponse: - # Errors convert per attempt, so the retry helper sees - # SandboxGatewayError instead of the generated ResponseError. + # Convert per attempt: the retry helper must see + # SandboxGatewayError, not the generated ResponseError. try: return await self._api.get_filesystem(path=path) except Exception as error: From 7fdeb1807758b6a771e78d3681886a7e1f558485 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Thu, 1 Oct 2026 11:58:56 -0700 Subject: [PATCH 09/18] Add the starter-image library listing to the sandbox surface library_images() on both clients lists the platform's starter images over the regenerated typed client, dropping hidden and coming-soon entries to match the console's create form. Named after the library_models consumer catalog, not Blaxel's hub. Co-Authored-By: Claude Code --- baseten/sandbox/__init__.py | 4 ++ baseten/sandbox/_client.py | 34 +++++++++++++++++ baseten/sandbox/_info.py | 65 +++++++++++++++++++++++++++++++ tests/sandbox/test_parity.py | 28 ++++++++++++++ tests/sandbox/test_surface.py | 72 +++++++++++++++++++++++++++++++++++ 5 files changed, 203 insertions(+) diff --git a/baseten/sandbox/__init__.py b/baseten/sandbox/__init__.py index 1aeb477..5ed9047 100644 --- a/baseten/sandbox/__init__.py +++ b/baseten/sandbox/__init__.py @@ -13,6 +13,8 @@ ) from baseten.sandbox._errors import SandboxApiError, SandboxGatewayError from baseten.sandbox._info import ( + LibraryImage, + LibraryImagePort, SandboxEnvValue, SandboxInfo, SandboxProcessInfo, @@ -33,6 +35,8 @@ "AsyncSandboxFileSystem", "AsyncSandboxProcess", "AsyncSandboxTokenProvider", + "LibraryImage", + "LibraryImagePort", "Sandbox", "SandboxApiError", "SandboxClient", diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py index e1a9781..ca781c6 100644 --- a/baseten/sandbox/_client.py +++ b/baseten/sandbox/_client.py @@ -20,8 +20,10 @@ ) from baseten.sandbox._errors import to_sandbox_api_error from baseten.sandbox._info import ( + LibraryImage, SandboxEnvValue, SandboxInfo, + library_image_from_api, sandbox_envs_to_api, sandbox_info_from_api, ) @@ -372,6 +374,22 @@ def delete(self, name: str) -> SandboxInfo: raise to_sandbox_api_error(error, "control") from error return sandbox_info_from_api(sandbox) + def library_images(self) -> Sequence[LibraryImage]: + """List the platform's starter images from the platform starter-image library. + + Hidden and coming-soon entries are dropped, matching what the + console's create form shows. + """ + try: + catalog = self._sandbox_api.list_sandbox_library_images() + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return [ + library_image_from_api(image) + for image in catalog.root + if not image.hidden and not image.coming_soon + ] + def close(self) -> None: """Close the client and the connections every sandbox shares.""" for http_client in self._sandbox_http_clients: @@ -649,6 +667,22 @@ async def delete(self, name: str) -> SandboxInfo: raise to_sandbox_api_error(error, "control") from error return sandbox_info_from_api(sandbox) + async def library_images(self) -> Sequence[LibraryImage]: + """List the platform's starter images from the platform starter-image library. + + Hidden and coming-soon entries are dropped, matching what the + console's create form shows. + """ + try: + catalog = await self._sandbox_api.list_sandbox_library_images() + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return [ + library_image_from_api(image) + for image in catalog.root + if not image.hidden and not image.coming_soon + ] + async def close(self) -> None: """Close the client and the connections every sandbox shares.""" for http_client in self._sandbox_http_clients: diff --git a/baseten/sandbox/_info.py b/baseten/sandbox/_info.py index bf6f1e9..91bc813 100644 --- a/baseten/sandbox/_info.py +++ b/baseten/sandbox/_info.py @@ -98,6 +98,71 @@ class SandboxProcessInfo: completed_at: datetime | None = None +@dataclass +class LibraryImagePort: + """One port a starter image's services listen on.""" + + name: str | None = None + target: int | None = None + """Port number inside the sandbox.""" + + protocol: str | None = None + + +@dataclass +class LibraryImage: + """One starter image from the platform's starter-image library. + + Available to any sandbox without building or pushing; pass + :attr:`image` as the image when creating a sandbox. + """ + + name: str + """Stable identifier of the starter image.""" + + image: str + """Image reference, including its tag.""" + + display_name: str | None = None + description: str | None = None + long_description: str | None = None + memory: int | None = None + """Default memory allocation in megabytes.""" + + categories: list[str] = field(default_factory=list) + tags: list[str] = field(default_factory=list) + ports: list[LibraryImagePort] = field(default_factory=list) + icon_url: str | None = None + project_url: str | None = None + """Project page for the image's stack.""" + + enterprise: bool = False + """Whether the image is gated to enterprise workspaces.""" + + +def library_image_from_api( + image: baseten.client.managementapi.SandboxLibraryImage, +) -> LibraryImage: + """Convert the control plane's starter-image catalog entry.""" + return LibraryImage( + name=image.name, + image=image.image, + display_name=image.displayName, + description=image.description, + long_description=image.longDescription, + memory=image.memory, + categories=image.categories or [], + tags=image.tags or [], + ports=[ + LibraryImagePort(name=port.name, target=port.target, protocol=port.protocol) + for port in image.ports or [] + ], + icon_url=image.icon, + project_url=image.url, + enterprise=bool(image.enterprise), + ) + + def sandbox_info_from_api(sandbox: baseten.client.managementapi.Sandbox) -> SandboxInfo: """Convert the control plane's sandbox record.""" return SandboxInfo( diff --git a/tests/sandbox/test_parity.py b/tests/sandbox/test_parity.py index 03b592f..ea565a5 100644 --- a/tests/sandbox/test_parity.py +++ b/tests/sandbox/test_parity.py @@ -35,6 +35,14 @@ # File-write fields left to the raw client. INTENTIONALLY_NOT_TRANSLATED_FILE = {"isDirectory", "permissions"} +# Starter-image fields not on the curated record: icon variants, plus the flags driving the library_images filter. +INTENTIONALLY_NOT_TRANSLATED_LIBRARY_IMAGE = { + "iconLight", + "iconDark", + "hidden", + "coming_soon", +} + def test_create_forwards_every_translated_field() -> None: forwarded = { @@ -116,3 +124,23 @@ def test_fs_write_forwards_every_translated_field() -> None: fields = set(baseten.client.sandboxapi.FileRequest.model_fields) assert forwarded <= fields assert fields - forwarded == INTENTIONALLY_NOT_TRANSLATED_FILE + + +def test_library_image_translation_covers_every_translated_field() -> None: + translated = { + "name", + "displayName", + "image", + "description", + "longDescription", + "memory", + "categories", + "tags", + "ports", + "icon", + "url", + "enterprise", + } + fields = set(baseten.client.managementapi.SandboxLibraryImage.model_fields) + assert translated <= fields + assert fields - translated == INTENTIONALLY_NOT_TRANSLATED_LIBRARY_IMAGE diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py index 56dde33..b89b159 100644 --- a/tests/sandbox/test_surface.py +++ b/tests/sandbox/test_surface.py @@ -540,3 +540,75 @@ def test_typed_error_body_becomes_a_dict() -> None: assert exc_info.value.code == "NOT_FOUND" assert isinstance(exc_info.value.body, dict) client.close() + + +def library_catalog_response(request: httpx.Request) -> httpx.Response: + # Callable route: a bare list payload would be consumed one entry per request. + return httpx.Response( + 200, json=library_catalog(), headers={"content-type": "application/json"} + ) + + +def library_catalog() -> list[dict[str, Any]]: + return [ + { + "name": "expo", + "image": "blaxel/expo:latest", + "displayName": "Expo", + "description": "Expo app development", + "longDescription": "Build and run Expo apps.", + "memory": 4096, + "categories": ["web"], + "tags": ["react-native"], + "ports": [{"name": "http", "target": 3000, "protocol": "HTTP"}], + "icon": "https://images.b10.co/expo.svg", + "url": "https://expo.dev", + "enterprise": False, + "hidden": False, + "coming_soon": False, + }, + {"name": "hidden-one", "image": "blaxel/hidden:latest", "hidden": True}, + { + "name": "soon-one", + "image": "blaxel/soon:latest", + "hidden": False, + "coming_soon": True, + }, + ] + + +def test_library_images_filters_and_translates() -> None: + transport = RoutingTransport( + routes={"GET /v0/sandbox/hub": library_catalog_response} + ) + client = make_sync_client(transport) + + images = client.library_images() + + assert [image.name for image in images] == ["expo"] + image = images[0] + assert image.image == "blaxel/expo:latest" + assert image.display_name == "Expo" + assert image.memory == 4096 + assert image.categories == ["web"] + assert image.ports[0].target == 3000 + assert image.icon_url == "https://images.b10.co/expo.svg" + assert image.project_url == "https://expo.dev" + client.close() + + +@pytest.mark.asyncio +async def test_async_library_images_filters() -> None: + transport = RoutingTransport( + routes={"GET /v0/sandbox/hub": library_catalog_response} + ) + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + + images = await client.library_images() + + assert [image.name for image in images] == ["expo"] + await client.close() From 3fe0aae010d659eb2a6fdf9392446b1b80371584 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Thu, 1 Oct 2026 12:49:31 -0700 Subject: [PATCH 10/18] Add the starter-image library listing to the sandbox surface library_images() on both clients lists the platform's starter images over the served GET /v1/sandboxes/library_images, which drops hidden and coming-soon entries server-side. Named after the library_models consumer vocabulary, not Blaxel's hub. Co-Authored-By: Claude Code --- baseten/sandbox/__init__.py | 2 + baseten/sandbox/_client.py | 34 ++++++------- baseten/sandbox/_info.py | 29 +++++++++-- baseten/sandbox/_sandbox.py | 36 +++++++++++++- tests/sandbox/test_parity.py | 15 ++++-- tests/sandbox/test_surface.py | 94 +++++++++++++++++++++++++++-------- 6 files changed, 163 insertions(+), 47 deletions(-) diff --git a/baseten/sandbox/__init__.py b/baseten/sandbox/__init__.py index 5ed9047..557a28c 100644 --- a/baseten/sandbox/__init__.py +++ b/baseten/sandbox/__init__.py @@ -18,6 +18,7 @@ SandboxEnvValue, SandboxInfo, SandboxProcessInfo, + SandboxProcessLogs, ) from baseten.sandbox._retry import SandboxRetryOptions from baseten.sandbox._sandbox import ( @@ -47,6 +48,7 @@ "SandboxInfo", "SandboxProcess", "SandboxProcessInfo", + "SandboxProcessLogs", "SandboxRetryOptions", "SandboxTokenProvider", ] diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py index ca781c6..e9e4449 100644 --- a/baseten/sandbox/_client.py +++ b/baseten/sandbox/_client.py @@ -375,20 +375,19 @@ def delete(self, name: str) -> SandboxInfo: return sandbox_info_from_api(sandbox) def library_images(self) -> Sequence[LibraryImage]: - """List the platform's starter images from the platform starter-image library. + """List the platform's starter images from the starter-image library. - Hidden and coming-soon entries are dropped, matching what the - console's create form shows. + The server drops hidden and coming-soon entries. """ try: - catalog = self._sandbox_api.list_sandbox_library_images() + catalog = self._sandbox_api.list_sandbox_library_images( + params=baseten.client.managementapi.ListSandboxLibraryImagesParams( + team_id=self._options.team_id + ) + ) except Exception as error: raise to_sandbox_api_error(error, "control") from error - return [ - library_image_from_api(image) - for image in catalog.root - if not image.hidden and not image.coming_soon - ] + return [library_image_from_api(image) for image in catalog.items] def close(self) -> None: """Close the client and the connections every sandbox shares.""" @@ -668,20 +667,19 @@ async def delete(self, name: str) -> SandboxInfo: return sandbox_info_from_api(sandbox) async def library_images(self) -> Sequence[LibraryImage]: - """List the platform's starter images from the platform starter-image library. + """List the platform's starter images from the starter-image library. - Hidden and coming-soon entries are dropped, matching what the - console's create form shows. + The server drops hidden and coming-soon entries. """ try: - catalog = await self._sandbox_api.list_sandbox_library_images() + catalog = await self._sandbox_api.list_sandbox_library_images( + params=baseten.client.managementapi.ListSandboxLibraryImagesParams( + team_id=self._options.team_id + ) + ) except Exception as error: raise to_sandbox_api_error(error, "control") from error - return [ - library_image_from_api(image) - for image in catalog.root - if not image.hidden and not image.coming_soon - ] + return [library_image_from_api(image) for image in catalog.items] async def close(self) -> None: """Close the client and the connections every sandbox shares.""" diff --git a/baseten/sandbox/_info.py b/baseten/sandbox/_info.py index 91bc813..7fa52cf 100644 --- a/baseten/sandbox/_info.py +++ b/baseten/sandbox/_info.py @@ -98,6 +98,16 @@ class SandboxProcessInfo: completed_at: datetime | None = None +@dataclass +class SandboxProcessLogs: + """A process's captured output.""" + + stdout: str + stderr: str + logs: str + """Standard output and standard error, interleaved.""" + + @dataclass class LibraryImagePort: """One port a starter image's services listen on.""" @@ -147,15 +157,19 @@ def library_image_from_api( return LibraryImage( name=image.name, image=image.image, - display_name=image.displayName, + display_name=image.display_name, description=image.description, - long_description=image.longDescription, + long_description=image.long_description, memory=image.memory, categories=image.categories or [], tags=image.tags or [], ports=[ - LibraryImagePort(name=port.name, target=port.target, protocol=port.protocol) - for port in image.ports or [] + LibraryImagePort( + name=port.name, + target=port.target, + protocol=str(port.protocol) if port.protocol else None, + ) + for port in (image.ports.root if image.ports else []) ], icon_url=image.icon, project_url=image.url, @@ -224,6 +238,13 @@ def process_info_from_api( ) +def process_logs_from_api( + logs: baseten.client.sandboxapi.ProcessLogs, +) -> SandboxProcessLogs: + """Convert the execution plane's captured process output.""" + return SandboxProcessLogs(stdout=logs.stdout, stderr=logs.stderr, logs=logs.logs) + + def parse_exec_timestamp(value: str) -> datetime: """Parse an execution-plane timestamp, ISO 8601 or RFC 1123.""" try: diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py index 8cbab4e..b3a318a 100644 --- a/baseten/sandbox/_sandbox.py +++ b/baseten/sandbox/_sandbox.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Self import httpx @@ -20,7 +20,9 @@ from baseten.sandbox._info import ( SandboxInfo, SandboxProcessInfo, + SandboxProcessLogs, process_info_from_api, + process_logs_from_api, ) from baseten.sandbox._retry import ( SandboxRetryOptions, @@ -473,6 +475,22 @@ def exec( raise to_sandbox_api_error(error, "exec") from error return process_info_from_api(response) + def list(self) -> Sequence[SandboxProcessInfo]: + """List the sandbox's processes.""" + try: + response = self._api.get_process() + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + return [process_info_from_api(process) for process in response.root] + + def logs(self, identifier: str) -> SandboxProcessLogs: + """Get one process's captured output, by pid or name.""" + try: + response = self._api.get_process_logs(identifier=identifier) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + return process_logs_from_api(response) + class AsyncSandboxProcess: """Processes in a sandbox. Async variant.""" @@ -509,3 +527,19 @@ async def exec( except Exception as error: raise to_sandbox_api_error(error, "exec") from error return process_info_from_api(response) + + async def list(self) -> Sequence[SandboxProcessInfo]: + """List the sandbox's processes.""" + try: + response = await self._api.get_process() + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + return [process_info_from_api(process) for process in response.root] + + async def logs(self, identifier: str) -> SandboxProcessLogs: + """Get one process's captured output, by pid or name.""" + try: + response = await self._api.get_process_logs(identifier=identifier) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + return process_logs_from_api(response) diff --git a/tests/sandbox/test_parity.py b/tests/sandbox/test_parity.py index ea565a5..aa42f28 100644 --- a/tests/sandbox/test_parity.py +++ b/tests/sandbox/test_parity.py @@ -37,10 +37,11 @@ # Starter-image fields not on the curated record: icon variants, plus the flags driving the library_images filter. INTENTIONALLY_NOT_TRANSLATED_LIBRARY_IMAGE = { - "iconLight", - "iconDark", + "icon_light", + "icon_dark", "hidden", "coming_soon", + "creation_options", } @@ -119,6 +120,12 @@ def test_process_record_translation_covers_every_translated_field() -> None: assert fields - translated == INTENTIONALLY_NOT_TRANSLATED_PROCESS_RECORD +def test_process_logs_translation_covers_every_field() -> None: + translated = {"logs", "stderr", "stdout"} + fields = set(baseten.client.sandboxapi.ProcessLogs.model_fields) + assert fields == translated + + def test_fs_write_forwards_every_translated_field() -> None: forwarded = {"content"} fields = set(baseten.client.sandboxapi.FileRequest.model_fields) @@ -129,10 +136,10 @@ def test_fs_write_forwards_every_translated_field() -> None: def test_library_image_translation_covers_every_translated_field() -> None: translated = { "name", - "displayName", + "display_name", "image", "description", - "longDescription", + "long_description", "memory", "categories", "tags", diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py index b89b159..8d4bcaa 100644 --- a/tests/sandbox/test_surface.py +++ b/tests/sandbox/test_surface.py @@ -369,6 +369,66 @@ def test_fs_read_directory_raises() -> None: client.close() +def test_process_list_and_logs() -> None: + def process_list_response(request: httpx.Request) -> httpx.Response: + # Callable route: a bare list payload would be consumed one entry per request. + return httpx.Response( + 200, json=[process_record()], headers={"content-type": "application/json"} + ) + + transport = RoutingTransport( + routes={ + "GET /process": process_list_response, + "GET /process/123/logs": {"logs": "hi\n", "stdout": "hi\n", "stderr": ""}, + } + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + processes = sandbox.process.list() + logs = sandbox.process.logs("123") + + assert [process.pid for process in processes] == ["123"] + assert processes[0].status == "completed" + assert logs.logs == "hi\n" + assert logs.stderr == "" + client.close() + + +async def _async_sandbox(transport: RoutingTransport): + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + return client, await client.create(name="sb-1") + + +@pytest.mark.asyncio +async def test_async_process_list_and_logs() -> None: + def process_list_response(request: httpx.Request) -> httpx.Response: + # Callable route: a bare list payload would be consumed one entry per request. + return httpx.Response( + 200, json=[process_record()], headers={"content-type": "application/json"} + ) + + transport = RoutingTransport( + routes={ + "POST /v1/sandboxes/instances": created_response(), + "GET /process": process_list_response, + "GET /process/123/logs": {"logs": "hi\n", "stdout": "hi\n", "stderr": ""}, + } + ) + client, sandbox = await _async_sandbox(transport) + + processes = await sandbox.process.list() + logs = await sandbox.process.logs("123") + + assert [process.pid for process in processes] == ["123"] + assert logs.logs == "hi\n" + await client.close() + + def test_process_exec_returns_curated_info() -> None: transport = RoutingTransport(routes={}) client = make_full_client(transport) @@ -543,20 +603,23 @@ def test_typed_error_body_becomes_a_dict() -> None: def library_catalog_response(request: httpx.Request) -> httpx.Response: - # Callable route: a bare list payload would be consumed one entry per request. + # Callable route: a dict payload would be consumed as a scripted response. return httpx.Response( - 200, json=library_catalog(), headers={"content-type": "application/json"} + 200, + json={"items": library_catalog()}, + headers={"content-type": "application/json"}, ) def library_catalog() -> list[dict[str, Any]]: + # The server drops hidden and coming-soon entries before listing. return [ { "name": "expo", - "image": "blaxel/expo:latest", - "displayName": "Expo", + "image": "baseten/expo:latest", + "display_name": "Expo", "description": "Expo app development", - "longDescription": "Build and run Expo apps.", + "long_description": "Build and run Expo apps.", "memory": 4096, "categories": ["web"], "tags": ["react-native"], @@ -564,22 +627,13 @@ def library_catalog() -> list[dict[str, Any]]: "icon": "https://images.b10.co/expo.svg", "url": "https://expo.dev", "enterprise": False, - "hidden": False, - "coming_soon": False, - }, - {"name": "hidden-one", "image": "blaxel/hidden:latest", "hidden": True}, - { - "name": "soon-one", - "image": "blaxel/soon:latest", - "hidden": False, - "coming_soon": True, - }, + } ] -def test_library_images_filters_and_translates() -> None: +def test_library_images_translates_the_served_catalog() -> None: transport = RoutingTransport( - routes={"GET /v0/sandbox/hub": library_catalog_response} + routes={"GET /v1/sandboxes/library_images": library_catalog_response} ) client = make_sync_client(transport) @@ -587,7 +641,7 @@ def test_library_images_filters_and_translates() -> None: assert [image.name for image in images] == ["expo"] image = images[0] - assert image.image == "blaxel/expo:latest" + assert image.image == "baseten/expo:latest" assert image.display_name == "Expo" assert image.memory == 4096 assert image.categories == ["web"] @@ -598,9 +652,9 @@ def test_library_images_filters_and_translates() -> None: @pytest.mark.asyncio -async def test_async_library_images_filters() -> None: +async def test_async_library_images_translates() -> None: transport = RoutingTransport( - routes={"GET /v0/sandbox/hub": library_catalog_response} + routes={"GET /v1/sandboxes/library_images": library_catalog_response} ) client = AsyncSandboxClient( api_key="test-key", From 7502c7fc7982d82d8b3250a2ceb45fc1f4d890ef Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Fri, 2 Oct 2026 10:50:09 -0700 Subject: [PATCH 11/18] Follow the served record shape: state and display_name drop, create gains create_if_not_exists The curated SandboxInfo loses state and display_name with the spec forward-patch; create() drops display_name and gains the create_if_not_exists idempotent form the CLI already exposes. Co-Authored-By: Claude Code --- baseten/sandbox/_client.py | 14 +++++++------- baseten/sandbox/_info.py | 8 -------- tests/sandbox/test_parity.py | 4 +--- 3 files changed, 8 insertions(+), 18 deletions(-) diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py index e9e4449..5ab0c2e 100644 --- a/baseten/sandbox/_client.py +++ b/baseten/sandbox/_client.py @@ -103,8 +103,8 @@ def _sandbox_create_request( region: str | None, envs: list[baseten.client.managementapi.SandboxEnv] | None, labels: dict[str, str] | None, - display_name: str | None, external_id: str | None, + create_if_not_exists: bool, ) -> baseten.client.managementapi.CreateSandboxRequest: """Build the create body from only the fields the caller set. @@ -125,10 +125,10 @@ def _sandbox_create_request( create_fields["envs"] = envs if labels is not None: create_fields["labels"] = labels - if display_name is not None: - create_fields["display_name"] = display_name if external_id is not None: create_fields["external_id"] = external_id + if create_if_not_exists: + create_fields["create_if_not_exists"] = True return baseten.client.managementapi.CreateSandboxRequest(**create_fields) @@ -266,8 +266,8 @@ def create( region: str | None = None, envs: dict[str, SandboxEnvValue] | None = None, labels: dict[str, str] | None = None, - display_name: str | None = None, external_id: str | None = None, + create_if_not_exists: bool = False, ) -> Sandbox: """Create a sandbox. @@ -286,8 +286,8 @@ def create( region=region, envs=sandbox_envs_to_api(envs) if envs is not None else None, labels=labels, - display_name=display_name, external_id=external_id, + create_if_not_exists=create_if_not_exists, ), ) except Exception as error: @@ -557,8 +557,8 @@ async def create( region: str | None = None, envs: dict[str, SandboxEnvValue] | None = None, labels: dict[str, str] | None = None, - display_name: str | None = None, external_id: str | None = None, + create_if_not_exists: bool = False, ) -> AsyncSandbox: """Create a sandbox. @@ -577,8 +577,8 @@ async def create( region=region, envs=sandbox_envs_to_api(envs) if envs is not None else None, labels=labels, - display_name=display_name, external_id=external_id, + create_if_not_exists=create_if_not_exists, ), ) except Exception as error: diff --git a/baseten/sandbox/_info.py b/baseten/sandbox/_info.py index 7fa52cf..6cf35fe 100644 --- a/baseten/sandbox/_info.py +++ b/baseten/sandbox/_info.py @@ -24,14 +24,10 @@ "UPLOADING", ) -# Whether a deployed sandbox is running or idle in standby. Kept open. -SANDBOX_STATES = ("RUNNING", "STANDBY") - # Status of a process in a sandbox. Kept open. SANDBOX_PROCESS_STATUSES = ("running", "completed", "failed", "killed", "stopped") SandboxStatus = str -SandboxState = str SandboxProcessStatus = str @@ -59,7 +55,6 @@ class SandboxInfo: url: str | None = None """Base URL of the sandbox's execution API, once it has one.""" - state: SandboxState | None = None image: str | None = None """Image reference, including its tag.""" @@ -67,7 +62,6 @@ class SandboxInfo: """Memory in megabytes, which also sets the CPU allocation.""" region: str | None = None - display_name: str | None = None external_id: str | None = None """Caller-owned identifier for external lookups.""" @@ -183,14 +177,12 @@ def sandbox_info_from_api(sandbox: baseten.client.managementapi.Sandbox) -> Sand name=sandbox.name, url=str(sandbox.url) if sandbox.url is not None else None, status=str(sandbox.status), - state=str(sandbox.state) if sandbox.state is not None else None, image=sandbox.image, memory=sandbox.memory, region=sandbox.region, enabled=sandbox.enabled, envs=_envs_from_api(sandbox.envs), labels=dict(sandbox.labels.root) if sandbox.labels is not None else {}, - display_name=sandbox.display_name, external_id=sandbox.external_id, created_at=sandbox.created_at, updated_at=sandbox.updated_at, diff --git a/tests/sandbox/test_parity.py b/tests/sandbox/test_parity.py index aa42f28..c0c09d8 100644 --- a/tests/sandbox/test_parity.py +++ b/tests/sandbox/test_parity.py @@ -53,8 +53,8 @@ def test_create_forwards_every_translated_field() -> None: "region", "envs", "labels", - "display_name", "external_id", + "create_if_not_exists", } fields = set(baseten.client.managementapi.CreateSandboxRequest.model_fields) assert forwarded <= fields @@ -66,14 +66,12 @@ def test_record_translation_covers_every_translated_field() -> None: "name", "url", "status", - "state", "image", "memory", "region", "enabled", "envs", "labels", - "display_name", "external_id", "created_at", "updated_at", From 8f378b18f041d46eab386281106bbeee2e9ae20b Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Mon, 5 Oct 2026 11:32:09 -0700 Subject: [PATCH 12/18] Catch up with the Go and JS SDK learnings - sandbox_from_url on both clients; the standalone Sandbox constructors are gone, so a sandbox always shares its client's authentication, connections, and retry budgets. - get_info takes show_secrets. - A revoked token's later resends wait out the revocation's rounded-up cutoff second; one token exchange is bounded by 30s. - create fails loudly when the served record carries no URL. - Env values default to secret, matching the server's default; the old False default would have unmasked values. - LibraryImage flattens the catalog's creation options, matching Go. Co-Authored-By: Claude Code --- baseten/sandbox/_auth.py | 12 +++ baseten/sandbox/_client.py | 82 ++++++++++++----- baseten/sandbox/_info.py | 45 +++++++++- baseten/sandbox/_sandbox.py | 164 ++++++---------------------------- tests/sandbox/test_parity.py | 2 +- tests/sandbox/test_surface.py | 60 ++++++++++++- 6 files changed, 202 insertions(+), 163 deletions(-) diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py index 1403db7..c090ecb 100644 --- a/baseten/sandbox/_auth.py +++ b/baseten/sandbox/_auth.py @@ -2,6 +2,7 @@ import asyncio import threading +import time from collections.abc import Awaitable, Callable from datetime import UTC, datetime, timedelta from typing import Protocol @@ -15,6 +16,13 @@ # one; beyond that, rejection points at something a new token cannot fix. _TOKEN_INVALIDATION_MAX_RETRIES = 2 +# A revocation rejects every token issued before its cutoff, the event's time +# rounded up to the next whole second. A token minted right after the event +# can fall before the cutoff and be rejected too, so after the first resend, +# each one waits this long to be minted past it. A module variable so tests +# can shorten it. +_TOKEN_REVOKED_RETRY_DELAY_SECONDS = 1.0 + # Both planes send this header on revocation, which can precede the stated # expiry. The request was untouched, so re-sending with a new token is safe # even when it has side effects. @@ -208,6 +216,8 @@ def handle_request(self, request: httpx.Request) -> httpx.Response: return response self._token_source.invalidate(token) response.close() + if attempt > 0: + time.sleep(_TOKEN_REVOKED_RETRY_DELAY_SECONDS) raise AssertionError # pragma: no cover - loop returns first def close(self) -> None: @@ -241,6 +251,8 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: return response await self._token_source.invalidate(token) await response.aclose() + if attempt > 0: + await asyncio.sleep(_TOKEN_REVOKED_RETRY_DELAY_SECONDS) raise AssertionError # pragma: no cover - loop returns first async def aclose(self) -> None: diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py index 5ab0c2e..e374629 100644 --- a/baseten/sandbox/_client.py +++ b/baseten/sandbox/_client.py @@ -36,6 +36,10 @@ # needs. Callers pass a stricter `timeout` when wanted. DEFAULT_TIMEOUT = httpx.Timeout(None, connect=10.0) +# One token exchange is bounded on its own, so a stalled mint is not what +# every caller waits on. A module variable so tests can shorten it. +_TOKEN_MINT_TIMEOUT_SECONDS = 30.0 + @dataclass(frozen=True) class SandboxClientOptions: @@ -192,7 +196,7 @@ def __init__( self._options, f"Bearer {api_key}" if api_key != "" else None, ), - timeout=self._timeout, + timeout=httpx.Timeout(_TOKEN_MINT_TIMEOUT_SECONDS), ) self._token_source = SyncTokenSource( api_key=api_key, @@ -293,15 +297,22 @@ def create( except Exception as error: raise to_sandbox_api_error(error, "control") from error info = sandbox_info_from_api(sandbox) - return self._sandbox_for(info, info) + if info.url is None: + raise ValueError(f"sandbox {info.name} was created but has no URL yet") + return self._sandbox_for(info.name, info.url, info) + + def get_info(self, name: str, *, show_secrets: bool = False) -> SandboxInfo: + """Get a sandbox's current record. - def get_info(self, name: str) -> SandboxInfo: - """Get a sandbox's current record.""" + ``show_secrets`` returns environment variable values unmasked; + it needs the workspace administrator role, and callers without it + still see masked values. + """ try: sandbox = self._sandbox_api.get_sandbox( sandbox_name=name, params=baseten.client.managementapi.GetSandboxParams( - team_id=self._options.team_id + team_id=self._options.team_id, show_secrets=show_secrets ), ) except Exception as error: @@ -315,7 +326,17 @@ def get(self, sandbox: str | SandboxInfo) -> Sandbox: record is already in hand. """ info = self.get_info(sandbox) if isinstance(sandbox, str) else sandbox - return self._sandbox_for(info) + if info.url is None: + raise ValueError(f"sandbox {info.name} has no URL yet") + return self._sandbox_for(info.name, info.url) + + def sandbox_from_url(self, url: str) -> Sandbox: + """Get a :class:`Sandbox` for a sandbox whose URL is already known. + + Makes no call. The sandbox shares this client's authentication, + connections, and retry budgets, and closes with it. + """ + return self._sandbox_for("", url) def list( self, @@ -405,20 +426,18 @@ def __exit__(self, *args: object) -> None: self.close() def _sandbox_for( - self, info: SandboxInfo, creation_info: SandboxInfo | None = None + self, name: str, url: str, creation_info: SandboxInfo | None = None ) -> Sandbox: - if info.url is None: - raise ValueError(f"sandbox {info.name} has no URL yet") http_client = httpx.Client( transport=self._auth_transport, - base_url=info.url, + base_url=url, headers=_request_headers(self._options, None), timeout=self._timeout, ) self._sandbox_http_clients.append(http_client) return Sandbox._shared( - name=info.name, - url=info.url, + name=name, + url=url, http_client=http_client, retries=self._options.retries, creation_info=creation_info, @@ -584,15 +603,22 @@ async def create( except Exception as error: raise to_sandbox_api_error(error, "control") from error info = sandbox_info_from_api(sandbox) - return await self._sandbox_for(info, info) + if info.url is None: + raise ValueError(f"sandbox {info.name} was created but has no URL yet") + return self._sandbox_for(info.name, info.url, info) + + async def get_info(self, name: str, *, show_secrets: bool = False) -> SandboxInfo: + """Get a sandbox's current record. - async def get_info(self, name: str) -> SandboxInfo: - """Get a sandbox's current record.""" + ``show_secrets`` returns environment variable values unmasked; + it needs the workspace administrator role, and callers without it + still see masked values. + """ try: sandbox = await self._sandbox_api.get_sandbox( sandbox_name=name, params=baseten.client.managementapi.GetSandboxParams( - team_id=self._options.team_id + team_id=self._options.team_id, show_secrets=show_secrets ), ) except Exception as error: @@ -606,7 +632,17 @@ async def get(self, sandbox: str | SandboxInfo) -> AsyncSandbox: record is already in hand. """ info = await self.get_info(sandbox) if isinstance(sandbox, str) else sandbox - return await self._sandbox_for(info) + if info.url is None: + raise ValueError(f"sandbox {info.name} has no URL yet") + return self._sandbox_for(info.name, info.url) + + def sandbox_from_url(self, url: str) -> AsyncSandbox: + """Get an :class:`AsyncSandbox` for a sandbox whose URL is already known. + + Makes no call. The sandbox shares this client's authentication, + connections, and retry budgets, and closes with it. + """ + return self._sandbox_for("", url) async def list( self, @@ -696,21 +732,19 @@ async def __aenter__(self) -> Self: async def __aexit__(self, *args: object) -> None: await self.close() - async def _sandbox_for( - self, info: SandboxInfo, creation_info: SandboxInfo | None = None + def _sandbox_for( + self, name: str, url: str, creation_info: SandboxInfo | None = None ) -> AsyncSandbox: - if info.url is None: - raise ValueError(f"sandbox {info.name} has no URL yet") http_client = httpx.AsyncClient( transport=self._auth_transport, - base_url=info.url, + base_url=url, headers=_request_headers(self._options, None), timeout=self._timeout, ) self._sandbox_http_clients.append(http_client) return AsyncSandbox._shared( - name=info.name, - url=info.url, + name=name, + url=url, http_client=http_client, retries=self._options.retries, creation_info=creation_info, diff --git a/baseten/sandbox/_info.py b/baseten/sandbox/_info.py index 6cf35fe..2273eb8 100644 --- a/baseten/sandbox/_info.py +++ b/baseten/sandbox/_info.py @@ -36,7 +36,10 @@ class SandboxEnvValue: """Value of an environment variable in a sandbox.""" value: str - secret: bool = False + # The server treats every env value as secret by default, so the safe + # default here is secret too; False opts out and unmasks the value in + # responses. + secret: bool = True @dataclass @@ -143,6 +146,31 @@ class LibraryImage: enterprise: bool = False """Whether the image is gated to enterprise workspaces.""" + creation_extra_args: dict[str, str] = field(default_factory=dict) + """Kernel selection arguments suggested when creating a sandbox.""" + + creation_volumes: list[LibraryImageVolume] = field(default_factory=list) + """Volume attachments suggested when creating a sandbox.""" + + +@dataclass +class LibraryImageVolume: + """A volume attachment a starter image suggests.""" + + name: str | None = None + """Volume name, or an internal identifier for an ephemeral volume.""" + + mount_path: str | None = None + """Absolute filesystem path where the volume is mounted.""" + + type: str | None = None + """Volume type, empty for a persistent volume.""" + + size_mb: int | None = None + """Storage capacity in megabytes for an ephemeral volume.""" + + read_only: bool = False + def library_image_from_api( image: baseten.client.managementapi.SandboxLibraryImage, @@ -168,6 +196,21 @@ def library_image_from_api( icon_url=image.icon, project_url=image.url, enterprise=bool(image.enterprise), + creation_extra_args=dict(image.creation_options.extra_args or {}) + if image.creation_options + else {}, + creation_volumes=[ + LibraryImageVolume( + name=volume.name, + mount_path=volume.mount_path, + type=volume.type, + size_mb=volume.size_mb, + read_only=bool(volume.read_only), + ) + for volume in (image.creation_options.volumes or []) + ] + if image.creation_options + else [], ) diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py index b3a318a..5e00a08 100644 --- a/baseten/sandbox/_sandbox.py +++ b/baseten/sandbox/_sandbox.py @@ -1,21 +1,11 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from typing import Self import httpx import baseten.client.sandboxapi -from baseten.client._user_agent import with_user_agent -from baseten.sandbox._auth import ( - AsyncSandboxTokenProvider, - AsyncTokenSource, - SandboxTokenProvider, - SyncTokenSource, - _AsyncAuthTransport, - _SyncAuthTransport, - resolve_http2_flag, -) from baseten.sandbox._errors import to_sandbox_api_error from baseten.sandbox._info import ( SandboxInfo, @@ -38,66 +28,21 @@ class Sandbox: """One sandbox, reached directly at its own URL. - Get one from :meth:`baseten.sandbox.SandboxClient.create` or - :meth:`baseten.sandbox.SandboxClient.get`, or construct one directly when - the URL and a token provider are already known. A sandbox from a client - shares that client's connections and is closed with it; one constructed - directly owns its connections and :meth:`close` closes them. + Get one from :meth:`baseten.sandbox.SandboxClient.create`, + :meth:`baseten.sandbox.SandboxClient.get`, or + :meth:`baseten.sandbox.SandboxClient.sandbox_from_url`. A sandbox + always shares its client's authentication, connections, and retry + budgets, and closes with it. """ - def __init__( - self, - *, - name: str, - url: str, - token_provider: SandboxTokenProvider | None = None, - headers: Mapping[str, str] | None = None, - http2: bool | None = None, - timeout: httpx.Timeout | None = None, - retries: SandboxRetryOptions | None = None, - transport_override: httpx.BaseTransport | None = None, - ) -> None: - """Construct a standalone sandbox client. - - Args: - name: Name of the sandbox. - url: Base URL of the sandbox's execution API, as reported in - ``SandboxInfo.url``. - token_provider: Returns the bearer token for each request. When - unset, no Authorization is sent. - headers: Additional headers to send on every request. - http2: Use HTTP/2 when the optional h2 package allows. None - enables it when h2 is importable, True requires it. - timeout: HTTP timeouts. Defaults to bounding only connect time. - retries: Retry budgets for transient failures. - transport_override: HTTP transport for tests with a mock. - """ - self._name = name - self._url = url.rstrip("/") - self._retries = retries or SandboxRetryOptions() - self._pool = ( - transport_override - if transport_override is not None - else httpx.HTTPTransport(http2=resolve_http2_flag(http2)) - ) - self._token_source = SyncTokenSource( - api_key="", token_provider=token_provider, mint=None - ) - self._auth_transport = _SyncAuthTransport(self._pool, self._token_source) - request_headers = with_user_agent({**(headers or {})}) - self._http_client = httpx.Client( - transport=self._auth_transport, - base_url=self._url, - headers=request_headers, - timeout=timeout - if timeout is not None - else httpx.Timeout(None, connect=10.0), - ) - self._api = baseten.client.sandboxapi.ApiClient(self._http_client) - self._fs: SandboxFileSystem | None = None - self._process: SandboxProcess | None = None - self._info: SandboxInfo | None = None - self._owns_pool = True + _name: str + _url: str + _retries: SandboxRetryOptions + _http_client: httpx.Client + _api: baseten.client.sandboxapi.ApiClient + _fs: SandboxFileSystem | None + _process: SandboxProcess | None + _info: SandboxInfo | None @classmethod def _shared( @@ -118,7 +63,6 @@ def _shared( sandbox._fs = None sandbox._process = None sandbox._info = creation_info - sandbox._owns_pool = False return sandbox @property @@ -180,75 +124,31 @@ def network(self) -> object: """Network access to ports in the sandbox. Not implemented yet.""" raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) - def close(self) -> None: - """Close this sandbox's connections, when it owns them. - - A sandbox from a :class:`baseten.sandbox.SandboxClient` shares that - client's connections and is closed with it, so this does nothing. - """ - if self._owns_pool: - self._http_client.close() - def __enter__(self) -> Self: return self def __exit__(self, *args: object) -> None: - self.close() + pass class AsyncSandbox: """One sandbox, reached directly at its own URL. Async variant. - Get one from :meth:`baseten.sandbox.AsyncSandboxClient.create` or - :meth:`baseten.sandbox.AsyncSandboxClient.get`, or construct one directly - when the URL and a token provider are already known. + Get one from :meth:`baseten.sandbox.AsyncSandboxClient.create`, + :meth:`baseten.sandbox.AsyncSandboxClient.get`, or + :meth:`baseten.sandbox.AsyncSandboxClient.sandbox_from_url`. The sandbox + shares its client's authentication, connections, and retry budgets, and + closes with it. """ - def __init__( - self, - *, - name: str, - url: str, - token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None = None, - headers: Mapping[str, str] | None = None, - http2: bool | None = None, - timeout: httpx.Timeout | None = None, - retries: SandboxRetryOptions | None = None, - transport_override: httpx.AsyncBaseTransport | None = None, - ) -> None: - """Construct a standalone async sandbox client. - - The token provider may return a string or an awaitable of one. Other - args match :class:`Sandbox`. - """ - self._name = name - self._url = url.rstrip("/") - self._retries = retries or SandboxRetryOptions() - self._pool = ( - transport_override - if transport_override is not None - else httpx.AsyncHTTPTransport(http2=resolve_http2_flag(http2)) - ) - self._token_source = AsyncTokenSource( - api_key="", - token_provider=token_provider, - mint=None, # type: ignore[arg-type] - ) - self._auth_transport = _AsyncAuthTransport(self._pool, self._token_source) - request_headers = with_user_agent({**(headers or {})}) - self._http_client = httpx.AsyncClient( - transport=self._auth_transport, - base_url=self._url, - headers=request_headers, - timeout=timeout - if timeout is not None - else httpx.Timeout(None, connect=10.0), - ) - self._api = baseten.client.sandboxapi.AsyncApiClient(self._http_client) - self._fs: AsyncSandboxFileSystem | None = None - self._process: AsyncSandboxProcess | None = None - self._info: SandboxInfo | None = None - self._owns_pool = True + _name: str + _url: str + _retries: SandboxRetryOptions + _http_client: httpx.AsyncClient + _api: baseten.client.sandboxapi.AsyncApiClient + _fs: AsyncSandboxFileSystem | None + _process: AsyncSandboxProcess | None + _info: SandboxInfo | None @classmethod def _shared( @@ -269,7 +169,6 @@ def _shared( sandbox._fs = None sandbox._process = None sandbox._info = creation_info - sandbox._owns_pool = False return sandbox @property @@ -331,16 +230,11 @@ def network(self) -> object: """Network access to ports in the sandbox. Not implemented yet.""" raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) - async def close(self) -> None: - """Close this sandbox's connections, when it owns them.""" - if self._owns_pool: - await self._http_client.aclose() - async def __aenter__(self) -> Self: return self async def __aexit__(self, *args: object) -> None: - await self.close() + pass class SandboxFileSystem: diff --git a/tests/sandbox/test_parity.py b/tests/sandbox/test_parity.py index c0c09d8..61a5ceb 100644 --- a/tests/sandbox/test_parity.py +++ b/tests/sandbox/test_parity.py @@ -41,7 +41,6 @@ "icon_dark", "hidden", "coming_soon", - "creation_options", } @@ -145,6 +144,7 @@ def test_library_image_translation_covers_every_translated_field() -> None: "icon", "url", "enterprise", + "creation_options", } fields = set(baseten.client.managementapi.SandboxLibraryImage.model_fields) assert translated <= fields diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py index 8d4bcaa..2abbb42 100644 --- a/tests/sandbox/test_surface.py +++ b/tests/sandbox/test_surface.py @@ -1,11 +1,13 @@ from __future__ import annotations import json +import types from typing import Any import httpx import pytest +import baseten.sandbox._auth from baseten.sandbox import ( AsyncSandboxClient, SandboxApiError, @@ -562,7 +564,18 @@ def test_fs_read_retries_gateway_errors() -> None: client.close() -def test_persistent_revocation_returns_last_response() -> None: +def test_persistent_revocation_returns_last_response(monkeypatch) -> None: + # The first resend is immediate; later ones wait out the revocation's + # rounded-up cutoff second, here shortened to keep the test fast. + delays: list[float] = [] + monkeypatch.setattr( + baseten.sandbox._auth, "_TOKEN_REVOKED_RETRY_DELAY_SECONDS", 0.01 + ) + monkeypatch.setattr( + baseten.sandbox._auth, + "time", + types.SimpleNamespace(sleep=lambda s: delays.append(s)), + ) transport = RoutingTransport( routes={ "GET /v1/sandboxes/instances/sb-1": lambda request: revoked_token_response( @@ -577,9 +590,52 @@ def test_persistent_revocation_returns_last_response() -> None: client.get_info("sb-1") # Exhaustion surfaces the final 401 as a converted error, and every - # minted token was tried. + # minted token was tried, the later resends past the cutoff. assert exc_info.value.status == 401 assert transport.mint_count == 3 + assert delays == [0.01] + client.close() + + +def test_get_info_show_secrets_reaches_the_query() -> None: + seen_queries: list[str] = [] + + def show_secrets_route(request: httpx.Request) -> httpx.Response: + seen_queries.append(str(request.url.params)) + return httpx.Response( + 200, json=sandbox_record(), headers={"content-type": "application/json"} + ) + + transport = RoutingTransport( + routes={"GET /v1/sandboxes/instances/sb-1": show_secrets_route} + ) + client = make_sync_client(transport) + + client.get_info("sb-1") + client.get_info("sb-1", show_secrets=True) + + assert "show_secrets=false" in seen_queries[0] + assert "show_secrets=true" in seen_queries[1] + client.close() + + +def test_sandbox_from_url_makes_no_call_and_shares_auth() -> None: + transport = RoutingTransport( + routes={ + "POST /process": process_record(), + } + ) + client = make_sync_client(transport) + + sandbox = client.sandbox_from_url("https://sb-1.invalid") + assert sandbox.name == "" + assert sandbox.url == "https://sb-1.invalid" + # No control-plane call was made to resolve it. + assert transport.count("GET", "/v1/sandboxes/instances/sb-1") == 0 + + # The sandbox shares the client's token: its exec call carries it. + sandbox.process.exec("echo hi", wait_for_completion=True) + assert transport.authorizations("/process") == ["Bearer tok-1"] client.close() From d28f92c0453e204936640bda02452737099e7a45 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Mon, 5 Oct 2026 13:29:05 -0700 Subject: [PATCH 13/18] Address Copilot review on the catch-up - The async token exchange carries the same 30s bound as sync. - Process list and logs retry on the read and gateway budgets, as idempotent execution-plane reads. - A same-content file write is idempotent, so fs write retries on the upload budgets. - LibraryImageVolume exports with the other curated library types. Co-Authored-By: Claude Code --- baseten/sandbox/__init__.py | 2 + baseten/sandbox/_client.py | 2 +- baseten/sandbox/_sandbox.py | 158 +++++++++++++++++++++++++++--------- 3 files changed, 121 insertions(+), 41 deletions(-) diff --git a/baseten/sandbox/__init__.py b/baseten/sandbox/__init__.py index 557a28c..0b9a618 100644 --- a/baseten/sandbox/__init__.py +++ b/baseten/sandbox/__init__.py @@ -15,6 +15,7 @@ from baseten.sandbox._info import ( LibraryImage, LibraryImagePort, + LibraryImageVolume, SandboxEnvValue, SandboxInfo, SandboxProcessInfo, @@ -38,6 +39,7 @@ "AsyncSandboxTokenProvider", "LibraryImage", "LibraryImagePort", + "LibraryImageVolume", "Sandbox", "SandboxApiError", "SandboxClient", diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py index e374629..8ee349b 100644 --- a/baseten/sandbox/_client.py +++ b/baseten/sandbox/_client.py @@ -498,7 +498,7 @@ def __init__( self._timeout = self._options.timeout or DEFAULT_TIMEOUT self._mint_http_client = httpx.AsyncClient( transport=self._pool, - timeout=self._timeout, + timeout=httpx.Timeout(_TOKEN_MINT_TIMEOUT_SECONDS), base_url=_management_base_url(self._options), headers=_request_headers( self._options, diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py index 5e00a08..ad7fba8 100644 --- a/baseten/sandbox/_sandbox.py +++ b/baseten/sandbox/_sandbox.py @@ -101,7 +101,7 @@ def fs(self) -> SandboxFileSystem: def process(self) -> SandboxProcess: """Processes in the sandbox.""" if self._process is None: - self._process = SandboxProcess(self._api) + self._process = SandboxProcess(self._api, self._retries) return self._process @property @@ -207,7 +207,7 @@ def fs(self) -> AsyncSandboxFileSystem: def process(self) -> AsyncSandboxProcess: """Processes in the sandbox.""" if self._process is None: - self._process = AsyncSandboxProcess(self._api) + self._process = AsyncSandboxProcess(self._api, self._retries) return self._process @property @@ -274,14 +274,26 @@ def attempt() -> baseten.client.sandboxapi.GetFilesystemResponse: raise ValueError(f"{path} is a directory, not a file") def write(self, path: str, content: str) -> None: - """Write a text file, creating it or replacing its content.""" - try: - self._api.put_filesystem( - path=path, - request=baseten.client.sandboxapi.FileRequest(content=content), - ) - except Exception as error: - raise to_sandbox_api_error(error, "exec") from error + """Write a text file, creating it or replacing its content. + + A write of the same content is idempotent, so it is retried on the + upload budgets after a dropped connection. + """ + + def attempt() -> None: + try: + self._api.put_filesystem( + path=path, + request=baseten.client.sandboxapi.FileRequest(content=content), + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + retry_idempotent( + attempt, + max_retries=self._retries.upload_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) class AsyncSandboxFileSystem: @@ -323,21 +335,38 @@ async def attempt() -> baseten.client.sandboxapi.GetFilesystemResponse: raise ValueError(f"{path} is a directory, not a file") async def write(self, path: str, content: str) -> None: - """Write a text file, creating it or replacing its content.""" - try: - await self._api.put_filesystem( - path=path, - request=baseten.client.sandboxapi.FileRequest(content=content), - ) - except Exception as error: - raise to_sandbox_api_error(error, "exec") from error + """Write a text file, creating it or replacing its content. + + A write of the same content is idempotent, so it is retried on the + upload budgets after a dropped connection. + """ + + async def attempt() -> None: + try: + await self._api.put_filesystem( + path=path, + request=baseten.client.sandboxapi.FileRequest(content=content), + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + await retry_idempotent_async( + attempt, + max_retries=self._retries.upload_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) class SandboxProcess: """Processes in a sandbox.""" - def __init__(self, api: baseten.client.sandboxapi.ApiClient) -> None: + def __init__( + self, + api: baseten.client.sandboxapi.ApiClient, + retries: SandboxRetryOptions, + ) -> None: self._api = api + self._retries = retries def exec( self, @@ -370,27 +399,54 @@ def exec( return process_info_from_api(response) def list(self) -> Sequence[SandboxProcessInfo]: - """List the sandbox's processes.""" - try: - response = self._api.get_process() - except Exception as error: - raise to_sandbox_api_error(error, "exec") from error + """List the sandbox's processes. + + Retried as an idempotent read, on the configured budgets. + """ + + def attempt() -> baseten.client.sandboxapi.GetProcessResponse: + try: + return self._api.get_process() + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + response = retry_idempotent( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) return [process_info_from_api(process) for process in response.root] def logs(self, identifier: str) -> SandboxProcessLogs: - """Get one process's captured output, by pid or name.""" - try: - response = self._api.get_process_logs(identifier=identifier) - except Exception as error: - raise to_sandbox_api_error(error, "exec") from error + """Get one process's captured output, by pid or name. + + Retried as an idempotent read, on the configured budgets. + """ + + def attempt() -> baseten.client.sandboxapi.ProcessLogs: + try: + return self._api.get_process_logs(identifier=identifier) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + response = retry_idempotent( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) return process_logs_from_api(response) class AsyncSandboxProcess: """Processes in a sandbox. Async variant.""" - def __init__(self, api: baseten.client.sandboxapi.AsyncApiClient) -> None: + def __init__( + self, + api: baseten.client.sandboxapi.AsyncApiClient, + retries: SandboxRetryOptions, + ) -> None: self._api = api + self._retries = retries async def exec( self, @@ -423,17 +479,39 @@ async def exec( return process_info_from_api(response) async def list(self) -> Sequence[SandboxProcessInfo]: - """List the sandbox's processes.""" - try: - response = await self._api.get_process() - except Exception as error: - raise to_sandbox_api_error(error, "exec") from error + """List the sandbox's processes. + + Retried as an idempotent read, on the configured budgets. + """ + + async def attempt() -> baseten.client.sandboxapi.GetProcessResponse: + try: + return await self._api.get_process() + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + response = await retry_idempotent_async( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) return [process_info_from_api(process) for process in response.root] async def logs(self, identifier: str) -> SandboxProcessLogs: - """Get one process's captured output, by pid or name.""" - try: - response = await self._api.get_process_logs(identifier=identifier) - except Exception as error: - raise to_sandbox_api_error(error, "exec") from error + """Get one process's captured output, by pid or name. + + Retried as an idempotent read, on the configured budgets. + """ + + async def attempt() -> baseten.client.sandboxapi.ProcessLogs: + try: + return await self._api.get_process_logs(identifier=identifier) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + response = await retry_idempotent_async( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) return process_logs_from_api(response) From e0fc7c664b29c6084d26dcf4e0853d4c70a14700 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Mon, 5 Oct 2026 14:10:52 -0700 Subject: [PATCH 14/18] Address Copilot's second review round - The async auth transport buffers with aread, since the request body can be an async stream. - SandboxApiError keeps an error-only body's error as its code and description; the gate checked the wrong field. - Sandbox drops its context-manager methods: a sandbox closes with its client, so the protocol advertised cleanup it does not own. - Async fs read and write retry paths covered by tests, mirroring the sync pair. Co-Authored-By: Claude Code --- baseten/sandbox/_auth.py | 4 ++- baseten/sandbox/_errors.py | 2 +- baseten/sandbox/_sandbox.py | 13 --------- tests/sandbox/test_surface.py | 52 +++++++++++++++++++++++++++++++++++ 4 files changed, 56 insertions(+), 15 deletions(-) diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py index c090ecb..804a2ac 100644 --- a/baseten/sandbox/_auth.py +++ b/baseten/sandbox/_auth.py @@ -234,7 +234,9 @@ def __init__( self._token_source = token_source async def handle_async_request(self, request: httpx.Request) -> httpx.Response: - request.read() + # aread: the body may be backed by an async stream, and a re-send + # replays it. + await request.aread() for attempt in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): token = await self._token_source.token() if token is not None: diff --git a/baseten/sandbox/_errors.py b/baseten/sandbox/_errors.py index 46cb8e5..9e918e2 100644 --- a/baseten/sandbox/_errors.py +++ b/baseten/sandbox/_errors.py @@ -25,7 +25,7 @@ def __init__(self, status: int, body: object) -> None: code = ( body_code if isinstance(body_code, str) - else (body_error if isinstance(body_message, str) else None) + else (body_error if isinstance(body_error, str) else None) ) description = None if isinstance(body_message, str): diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py index ad7fba8..30b8e60 100644 --- a/baseten/sandbox/_sandbox.py +++ b/baseten/sandbox/_sandbox.py @@ -1,7 +1,6 @@ from __future__ import annotations from collections.abc import Sequence -from typing import Self import httpx @@ -124,12 +123,6 @@ def network(self) -> object: """Network access to ports in the sandbox. Not implemented yet.""" raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) - def __enter__(self) -> Self: - return self - - def __exit__(self, *args: object) -> None: - pass - class AsyncSandbox: """One sandbox, reached directly at its own URL. Async variant. @@ -230,12 +223,6 @@ def network(self) -> object: """Network access to ports in the sandbox. Not implemented yet.""" raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) - async def __aenter__(self) -> Self: - return self - - async def __aexit__(self, *args: object) -> None: - pass - class SandboxFileSystem: """Files and directories in a sandbox.""" diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py index 2abbb42..a4d02e9 100644 --- a/tests/sandbox/test_surface.py +++ b/tests/sandbox/test_surface.py @@ -351,6 +351,58 @@ def test_fs_read_retries_dropped_connection() -> None: client.close() +@pytest.mark.asyncio +async def test_async_fs_read_and_write_retry_transient_failures() -> None: + transport = RoutingTransport( + routes={ + "POST /v1/sandboxes/instances": created_response(), + "GET /filesystem/file.txt": [ + httpx.ConnectError("dropped"), + httpx.Response(502, text="bad gateway"), + file_record(), + ], + "PUT /filesystem/file.txt": [ + httpx.ConnectError("dropped"), + httpx.Response(503, text="unavailable"), + {"message": "written"}, + ], + } + ) + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + sandbox = await client.create(name="sb-1") + + content = await sandbox.fs.read("file.txt") + await sandbox.fs.write("file.txt", content) + + assert content == "file text" + assert transport.count("GET", "/filesystem/file.txt") == 3 + assert transport.count("PUT", "/filesystem/file.txt") == 3 + await client.close() + + +def test_error_only_body_carries_its_error_as_code_and_description() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": httpx.Response( + 409, + json={"error": "SANDBOX_NAME_TAKEN"}, + ) + } + ) + client = make_sync_client(transport) + + with pytest.raises(SandboxApiError) as exc_info: + client.get_info("sb-1") + + assert exc_info.value.code == "SANDBOX_NAME_TAKEN" + assert "SANDBOX_NAME_TAKEN" in str(exc_info.value) + client.close() + + def test_fs_read_directory_raises() -> None: transport = RoutingTransport( routes={ From 060a671c99eada808e9c6623d1c06ae72aaf3191 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Mon, 5 Oct 2026 14:28:47 -0700 Subject: [PATCH 15/18] Send a stream-backed body once, without revocation resend Matching the Go rule: only a body that can be replayed is ever sent again. An in-memory body re-sends on a revoked token; one still backed by a stream, such as a raw multipart upload, is sent once, since replaying it would mean materializing it. The eager read/aread that buffered every body is gone. Co-Authored-By: Claude Code --- baseten/sandbox/_auth.py | 32 ++++++++++++++++++++++++-------- tests/sandbox/test_surface.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 8 deletions(-) diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py index 804a2ac..69ccb62 100644 --- a/baseten/sandbox/_auth.py +++ b/baseten/sandbox/_auth.py @@ -188,8 +188,9 @@ class _SyncAuthTransport(httpx.BaseTransport): """Adds a bearer token from the source to every request. A request whose token was revoked is re-sent with a fresh one, up to - :data:`_TOKEN_INVALIDATION_MAX_RETRIES` times. The request body is - buffered first so a re-send replays it. + :data:`_TOKEN_INVALIDATION_MAX_RETRIES` times. Only a body already in + memory is re-sent: one still backed by a stream is sent once, since + replaying it would mean materializing it here. """ def __init__( @@ -199,14 +200,22 @@ def __init__( self._token_source = token_source def handle_request(self, request: httpx.Request) -> httpx.Response: - request.read() + # A body still backed by a stream, such as a raw multipart upload, + # cannot be replayed, and materializing it here is the memory + # exhaustion to avoid: it is sent once, without revocation resend. + try: + request.content + replayable = True + except httpx.RequestNotRead: + replayable = False for attempt in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): token = self._token_source.token() if token is not None: request.headers["Authorization"] = f"Bearer {token}" response = self._inner.handle_request(request) revoked = ( - token is not None + replayable + and token is not None and response.status_code == 401 and response.headers.get("x-blaxel-error-code") == _TOKEN_REVOKED_CODE ) @@ -234,16 +243,23 @@ def __init__( self._token_source = token_source async def handle_async_request(self, request: httpx.Request) -> httpx.Response: - # aread: the body may be backed by an async stream, and a re-send - # replays it. - await request.aread() + # A body still backed by an async stream, such as a raw multipart + # upload, cannot be replayed, and materializing it here is the + # memory exhaustion to avoid: it is sent once, without revocation + # resend. + try: + request.content + replayable = True + except httpx.RequestNotRead: + replayable = False for attempt in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): token = await self._token_source.token() if token is not None: request.headers["Authorization"] = f"Bearer {token}" response = await self._inner.handle_async_request(request) revoked = ( - token is not None + replayable + and token is not None and response.status_code == 401 and response.headers.get("x-blaxel-error-code") == _TOKEN_REVOKED_CODE ) diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py index a4d02e9..d06130b 100644 --- a/tests/sandbox/test_surface.py +++ b/tests/sandbox/test_surface.py @@ -384,6 +384,39 @@ async def test_async_fs_read_and_write_retry_transient_failures() -> None: await client.close() +def test_streaming_body_is_sent_once_on_revocation() -> None: + # A stream-backed body cannot be replayed, so a revoked token surfaces + # instead of a re-send that would have to materialize the stream. + sends: list[bytes] = [] + + def revoked_route(request: httpx.Request) -> httpx.Response: + request.read() + sends.append(b"".join(request.stream)) + return revoked_token_response("revoked") + + transport = RoutingTransport( + routes={"POST /filesystem/big.bin": revoked_route}, + mint_tokens=["tok-1"], + ) + client = make_sync_client(transport) + + stream = httpx.Client( + transport=client._auth_transport, + base_url=MANAGEMENT_URL, + headers=_request_headers_for_test(), + ) + response = stream.post("/filesystem/big.bin", content=iter([b"part-1", b"part-2"])) + + assert response.status_code == 401 + assert sends == [b"part-1part-2"] + assert transport.mint_count == 1 + client.close() + + +def _request_headers_for_test() -> dict[str, str]: + return {"user-agent": "test"} + + def test_error_only_body_carries_its_error_as_code_and_description() -> None: transport = RoutingTransport( routes={ From b8f74815559c87e952b46abd34418fa3601c587f Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Mon, 5 Oct 2026 14:36:06 -0700 Subject: [PATCH 16/18] Satisfy the linters on the stream probe Co-Authored-By: Claude Code --- baseten/sandbox/_auth.py | 4 ++-- tests/sandbox/test_surface.py | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py index 69ccb62..e3edf69 100644 --- a/baseten/sandbox/_auth.py +++ b/baseten/sandbox/_auth.py @@ -204,7 +204,7 @@ def handle_request(self, request: httpx.Request) -> httpx.Response: # cannot be replayed, and materializing it here is the memory # exhaustion to avoid: it is sent once, without revocation resend. try: - request.content + _ = request.content replayable = True except httpx.RequestNotRead: replayable = False @@ -248,7 +248,7 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: # memory exhaustion to avoid: it is sent once, without revocation # resend. try: - request.content + _ = request.content replayable = True except httpx.RequestNotRead: replayable = False diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py index d06130b..75f5729 100644 --- a/tests/sandbox/test_surface.py +++ b/tests/sandbox/test_surface.py @@ -390,8 +390,7 @@ def test_streaming_body_is_sent_once_on_revocation() -> None: sends: list[bytes] = [] def revoked_route(request: httpx.Request) -> httpx.Response: - request.read() - sends.append(b"".join(request.stream)) + sends.append(request.read()) return revoked_token_response("revoked") transport = RoutingTransport( From 09566e6b88f352cb0d22f2a04561a72cba208000 Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Mon, 5 Oct 2026 14:47:48 -0700 Subject: [PATCH 17/18] Refuse direct construction of Sandbox objects With no __init__, Sandbox() built an uninitialized object that failed later on access; construction now names the client entry points. _shared bypasses it with __new__. Co-Authored-By: Claude Code --- baseten/sandbox/_sandbox.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py index 30b8e60..b614a28 100644 --- a/baseten/sandbox/_sandbox.py +++ b/baseten/sandbox/_sandbox.py @@ -43,6 +43,14 @@ class Sandbox: _process: SandboxProcess | None _info: SandboxInfo | None + def __init__(self) -> None: + # _shared builds instances; direct construction cannot produce a + # usable one, so refuse it rather than fail later on access. + raise TypeError( + "build a Sandbox through SandboxClient.create, .get, or " + ".sandbox_from_url" + ) + @classmethod def _shared( cls, @@ -143,6 +151,14 @@ class AsyncSandbox: _process: AsyncSandboxProcess | None _info: SandboxInfo | None + def __init__(self) -> None: + # _shared builds instances; direct construction cannot produce a + # usable one, so refuse it rather than fail later on access. + raise TypeError( + "build an AsyncSandbox through AsyncSandboxClient.create, .get, " + "or .sandbox_from_url" + ) + @classmethod def _shared( cls, From c1f0e4c56d5bcf8112e0c43c7fb7b69f7d17e60f Mon Sep 17 00:00:00 2001 From: "Marius F. Killinger" Date: Mon, 5 Oct 2026 15:00:58 -0700 Subject: [PATCH 18/18] Always invalidate a revoked token, gating only the resend A stream-backed body that meets a revocation is still sent only once, but its token now leaves the cache, so the next request mints instead of reusing the known-bad one; the exhausted replayable attempts do too. Co-Authored-By: Claude Code --- baseten/sandbox/_auth.py | 28 ++++++++++++++++++++-------- baseten/sandbox/_sandbox.py | 3 +-- tests/sandbox/test_surface.py | 7 +++++++ 3 files changed, 28 insertions(+), 10 deletions(-) diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py index e3edf69..b24e0d5 100644 --- a/baseten/sandbox/_auth.py +++ b/baseten/sandbox/_auth.py @@ -214,16 +214,22 @@ def handle_request(self, request: httpx.Request) -> httpx.Response: request.headers["Authorization"] = f"Bearer {token}" response = self._inner.handle_request(request) revoked = ( - replayable - and token is not None + token is not None and response.status_code == 401 and response.headers.get("x-blaxel-error-code") == _TOKEN_REVOKED_CODE ) + if revoked: + # A revoked token never stays cached, even when the request + # cannot be re-sent or the attempts are exhausted. + self._token_source.invalidate(token) # The last attempt's response is returned even when still # revoked, so the generated client raises a meaningful error. - if not revoked or attempt == _TOKEN_INVALIDATION_MAX_RETRIES: + if ( + not revoked + or not replayable + or attempt == _TOKEN_INVALIDATION_MAX_RETRIES + ): return response - self._token_source.invalidate(token) response.close() if attempt > 0: time.sleep(_TOKEN_REVOKED_RETRY_DELAY_SECONDS) @@ -258,16 +264,22 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: request.headers["Authorization"] = f"Bearer {token}" response = await self._inner.handle_async_request(request) revoked = ( - replayable - and token is not None + token is not None and response.status_code == 401 and response.headers.get("x-blaxel-error-code") == _TOKEN_REVOKED_CODE ) + if revoked: + # A revoked token never stays cached, even when the request + # cannot be re-sent or the attempts are exhausted. + await self._token_source.invalidate(token) # The last attempt's response is returned even when still # revoked, so the generated client raises a meaningful error. - if not revoked or attempt == _TOKEN_INVALIDATION_MAX_RETRIES: + if ( + not revoked + or not replayable + or attempt == _TOKEN_INVALIDATION_MAX_RETRIES + ): return response - await self._token_source.invalidate(token) await response.aclose() if attempt > 0: await asyncio.sleep(_TOKEN_REVOKED_RETRY_DELAY_SECONDS) diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py index b614a28..e0e3c7c 100644 --- a/baseten/sandbox/_sandbox.py +++ b/baseten/sandbox/_sandbox.py @@ -47,8 +47,7 @@ def __init__(self) -> None: # _shared builds instances; direct construction cannot produce a # usable one, so refuse it rather than fail later on access. raise TypeError( - "build a Sandbox through SandboxClient.create, .get, or " - ".sandbox_from_url" + "build a Sandbox through SandboxClient.create, .get, or .sandbox_from_url" ) @classmethod diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py index 75f5729..d3b532c 100644 --- a/tests/sandbox/test_surface.py +++ b/tests/sandbox/test_surface.py @@ -409,6 +409,13 @@ def revoked_route(request: httpx.Request) -> httpx.Response: assert response.status_code == 401 assert sends == [b"part-1part-2"] assert transport.mint_count == 1 + + # The revoked token left the cache: the next request mints a fresh one + # instead of reusing the known-bad token. + transport.routes["GET /v1/sandboxes/instances/sb-1"] = sandbox_record() + info = client.get_info("sb-1") + assert info.name == "sb-1" + assert transport.mint_count == 2 client.close()