"""Connect the SDK abstractions through a simple, extensible sandbox interface.
ConTree's stateless execution model differs from a classic stateful sandbox:
filesystem state passes between executions through images and snapshots. This
module assembles clients, sessions, stores, and policies to provide a
familiar sandbox workflow without requiring users to connect each component.
The default lazy session reuses a VM between commands. Snapshots preserve files,
but do not preserve running processes or memory. Advanced users can select or
extend session classes, supply components, and configure policies while keeping
control over execution, state, and resource ownership.
"""
from __future__ import annotations
import asyncio
import sys
from collections.abc import Callable, Mapping
from contextlib import AbstractAsyncContextManager, AbstractContextManager, AsyncExitStack, ExitStack, suppress
from dataclasses import dataclass, field
from pathlib import Path
from types import TracebackType
from typing import Any, Generic, Literal, NoReturn, TypeVar, cast, overload
from contree_client.profiles import Profile
from contree_client.types import ContreeAsyncClient, ContreeSyncClient
from contree_sdk.connection import ConnectionConfig
from contree_sdk.connection import profile_from_args as profile_from_args # noqa: PLC0414 - public re-export
from contree_sdk.execution import AsyncExecutor, SyncExecutor
from contree_sdk.files import AsyncClientFileTransfer, AsyncFileTransfer, ClientFileTransfer, SyncFileTransfer
from contree_sdk.scope import StoreScope
from contree_sdk.session.asyncio import ContreeAsyncSession
from contree_sdk.session.cleanup import await_cleanup
from contree_sdk.session.commit_policy import AbstractCommitPolicy
from contree_sdk.session.lazy_async import AsyncLazySession
from contree_sdk.session.lazy_sync import LazySession
from contree_sdk.session.sync import ContreeSession
from contree_sdk.store import AsyncMemoryStore, AsyncStore, SyncMemoryStore, SyncStore
from contree_sdk.store.ttl import expiration
from contree_sdk.upload_cache import AsyncCachedFileTransfer, CachedFileTransfer
SyncSessionT = TypeVar("SyncSessionT", bound=SyncExecutor)
AsyncSessionT = TypeVar("AsyncSessionT", bound=AsyncExecutor)
ContextT = TypeVar("ContextT")
Client = ContreeSyncClient | ContreeAsyncClient
[docs]
@dataclass
class SandboxConfig:
image: str | None
url: str | None = None
project: str | None = None
token: str | None = field(default=None, repr=False)
auth_type: Literal["iam", "jwt"] | None = None
profile: Profile | str | None = None
config_path: str | Path | None = None
session_type: type[SyncExecutor | AsyncExecutor] | None = None
session_id: str | None = None
node_id: str | None = None
store: SyncStore | AsyncStore | None = None
cache_ttl: float = 60
cwd: str | None = None
env: dict[str, str] | None = None
commit_policy: AbstractCommitPolicy | None = None
file_transfer: SyncFileTransfer | AsyncFileTransfer | None = None
session_kwargs: dict[str, Any] = field(default_factory=dict)
client: Client | None = None
client_type: type[ContreeSyncClient | ContreeAsyncClient] | None = None
client_kwargs: dict[str, Any] = field(default_factory=dict, repr=False)
[docs]
def connection(self) -> ConnectionConfig:
return ConnectionConfig(
url=self.url,
project=self.project,
token=self.token,
auth_type=self.auth_type,
profile=self.profile,
config_path=self.config_path,
client=self.client,
client_type=self.client_type,
client_kwargs=self.client_kwargs,
)
[docs]
def resolve_profile(self) -> Profile:
return self.connection().resolve_profile()
[docs]
def validate(self, *, asynchronous: bool) -> None:
mode = "async with" if asynchronous else "with"
if self.session_type is not None and (
not isinstance(self.session_type, type)
or not issubclass(self.session_type, AsyncExecutor if asynchronous else SyncExecutor)
):
raise TypeError(f"session_type must be a compatible session class for {mode}")
self.connection().validate(asynchronous=asynchronous)
if self.store is not None and not isinstance(self.store, AsyncStore if asynchronous else SyncStore):
raise TypeError(f"store must implement the store interface for {mode}")
if self.file_transfer is not None and not isinstance(
self.file_transfer, AsyncFileTransfer if asynchronous else SyncFileTransfer
):
raise TypeError(f"file_transfer must implement the file transfer interface for {mode}")
if self.image is None and self.session_id is None:
raise ValueError("either image or session_id must be provided")
expiration(self.cache_ttl, 0)
[docs]
def session_arguments(
self, store: SyncStore | AsyncStore, transfer: SyncFileTransfer | AsyncFileTransfer
) -> tuple[dict[str, Any], dict[str, Any]]:
arguments = {
"image": self.image,
"session_id": self.session_id,
"node_id": self.node_id,
"store": store,
"file_transfer": transfer,
"cwd": self.cwd,
"env": self.env,
"commit_policy": self.commit_policy,
}
kwargs = self.session_kwargs.copy()
kwargs.pop("client", None)
kwargs.pop("session", None)
for name in arguments:
kwargs.pop(name, None)
return arguments, kwargs
[docs]
class SyncSandboxSessionContext(Generic[SyncSessionT]):
"""Construct a synchronous session and unwind its owned resources."""
def __init__(self, config: SandboxConfig) -> None:
self.config = config
self.stack = ExitStack()
self._entered = False
def _create(self) -> SyncExecutor:
config, stack = self.config, self.stack
client = config.connection().open_sync(stack)
store = cast("SyncStore | None", config.store)
if store is None:
store = SyncMemoryStore()
stack.callback(store.close)
transfer = cast("SyncFileTransfer | None", config.file_transfer)
if transfer is None:
transfer = ClientFileTransfer(client)
transfer = CachedFileTransfer(transfer, store, scope=StoreScope.from_client(client), ttl=config.cache_ttl)
session_type = cast("type[SyncExecutor]", config.session_type or LazySession)
arguments, kwargs = config.session_arguments(store, transfer)
if issubclass(session_type, LazySession):
session = session_type(ContreeSession(client, **arguments), **kwargs)
else:
constructor = cast("Callable[..., SyncExecutor]", session_type)
session = constructor(client, **arguments, **kwargs)
if isinstance(session, AbstractContextManager):
# Register cleanup before a custom context's entry can fail.
stack.push(session)
session.__enter__() # noqa: PLC2801 - exit is registered before entry can fail
return session
[docs]
def __enter__(self) -> SyncSessionT:
if self._entered:
raise RuntimeError("a sandbox context can only be entered once")
self._entered = True
self.config.validate(asynchronous=False)
try:
return cast("SyncSessionT", self._create())
except BaseException:
with suppress(Exception):
self.stack.__exit__(*sys.exc_info())
raise
def __exit__(
self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
) -> None:
if exc_type is None:
self.stack.close()
else:
with suppress(Exception):
self.stack.__exit__(exc_type, exc, tb)
[docs]
class AsyncSandboxSessionContext(Generic[AsyncSessionT]):
"""Construct a native async session and join cleanup despite cancellation."""
def __init__(self, config: SandboxConfig) -> None:
self.config = config
self.stack = AsyncExitStack()
self._entered = False
async def _create(self) -> AsyncExecutor:
config, stack = self.config, self.stack
client = await config.connection().open_async(stack)
store = cast("AsyncStore | None", config.store)
if store is None:
store = AsyncMemoryStore()
stack.push_async_callback(store.close)
transfer = cast("AsyncFileTransfer | None", config.file_transfer)
if transfer is None:
transfer = AsyncClientFileTransfer(client)
transfer = AsyncCachedFileTransfer(transfer, store, scope=StoreScope.from_client(client), ttl=config.cache_ttl)
session_type = cast("type[AsyncExecutor]", config.session_type or AsyncLazySession)
arguments, kwargs = config.session_arguments(store, transfer)
if issubclass(session_type, AsyncLazySession):
base = ContreeAsyncSession(client, **arguments)
session = session_type(base, **kwargs)
else:
constructor = cast("Callable[..., AsyncExecutor]", session_type)
session = constructor(client, **arguments, **kwargs)
base = session if isinstance(session, ContreeAsyncSession) else None
if isinstance(session, AbstractAsyncContextManager):
# Register cleanup before a custom context's entry can fail.
stack.push_async_exit(session)
await session.__aenter__() # noqa: PLC2801 - exit is registered before entry can fail
if base is not None:
await base.ensure_ready()
return session
[docs]
async def __aenter__(self) -> AsyncSessionT:
if self._entered:
raise RuntimeError("a sandbox context can only be entered once")
self._entered = True
self.config.validate(asynchronous=True)
try:
return cast("AsyncSessionT", await self._create())
except BaseException:
with suppress(Exception, asyncio.CancelledError):
await await_cleanup(self._unwind(*sys.exc_info()))
raise
async def _unwind(
self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
) -> None:
await self.stack.__aexit__(exc_type, exc, tb)
async def __aexit__(
self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
) -> None:
cleanup = self._unwind(exc_type, exc, tb)
if exc_type is None:
await await_cleanup(cleanup)
else:
with suppress(Exception, asyncio.CancelledError):
await await_cleanup(cleanup)
[docs]
class SandboxSessionContext(Generic[SyncSessionT, AsyncSessionT]):
"""Compose sync and async contexts with one shared, single-use boundary.
Construction performs no I/O. Entry chooses the context mode explicitly.
Re-entry, nesting, and switching modes on the same object are rejected.
"""
def __init__(self, config: SandboxConfig) -> None:
self._sync = SyncSandboxSessionContext[SyncSessionT](config)
self._async = AsyncSandboxSessionContext[AsyncSessionT](config)
self._entered = False
def _begin(self) -> None:
if self._entered:
raise RuntimeError("a sandbox context can only be entered once")
self._entered = True
[docs]
def __enter__(self) -> SyncSessionT:
self._begin()
return self._sync.__enter__()
def __exit__(
self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
) -> None:
self._sync.__exit__(exc_type, exc, tb)
[docs]
async def __aenter__(self) -> AsyncSessionT:
self._begin()
return await self._async.__aenter__()
async def __aexit__(
self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
) -> None:
await self._async.__aexit__(exc_type, exc, tb)
@overload
def create_sandbox_session(
*,
session_type: None = None,
store: SyncStore | AsyncStore | None = None,
client: ContreeSyncClient | ContreeAsyncClient | None = None,
client_type: type[ContreeSyncClient | ContreeAsyncClient] | None = None,
file_transfer: SyncFileTransfer | AsyncFileTransfer | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> SandboxSessionContext[LazySession, AsyncLazySession]: ...
@overload
def create_sandbox_session(
*,
session_type: type[SyncSessionT],
store: SyncStore | None = None,
client: ContreeSyncClient | None = None,
client_type: type[ContreeSyncClient] | None = None,
file_transfer: SyncFileTransfer | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> SandboxSessionContext[SyncSessionT, NoReturn]: ...
@overload
def create_sandbox_session(
*,
session_type: type[AsyncSessionT],
store: AsyncStore | None = None,
client: ContreeAsyncClient | None = None,
client_type: type[ContreeAsyncClient] | None = None,
file_transfer: AsyncFileTransfer | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> SandboxSessionContext[NoReturn, AsyncSessionT]: ...
[docs]
def create_sandbox_session(
*,
session_type: type[SyncExecutor | AsyncExecutor] | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> SandboxSessionContext[Any, Any]:
"""Create a single-use context supporting both with and async with.
Pass constructor options for the selected class through session_kwargs.
Other keyword arguments configure the factory or its client.
Returns:
A composed context. See session_maker for all configuration parameters.
"""
return session_maker(SandboxSessionContext, session_type=session_type, session_kwargs=session_kwargs, **kwargs)
@overload
def create_sync_sandbox_session(
*,
session_type: None = None,
store: SyncStore | None = None,
client: ContreeSyncClient | None = None,
client_type: type[ContreeSyncClient] | None = None,
file_transfer: SyncFileTransfer | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> SyncSandboxSessionContext[LazySession]: ...
@overload
def create_sync_sandbox_session(
*,
session_type: type[SyncSessionT],
store: SyncStore | None = None,
client: ContreeSyncClient | None = None,
client_type: type[ContreeSyncClient] | None = None,
file_transfer: SyncFileTransfer | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> SyncSandboxSessionContext[SyncSessionT]: ...
[docs]
def create_sync_sandbox_session(
*,
session_type: type[SyncExecutor] | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> SyncSandboxSessionContext[Any]:
"""Create a sync-only context; its default session is LazySession.
Pass constructor options for the selected class through session_kwargs.
Other keyword arguments configure the factory or its client.
Returns:
A synchronous context with the configuration accepted by session_maker.
"""
return session_maker(SyncSandboxSessionContext, session_type=session_type, session_kwargs=session_kwargs, **kwargs)
@overload
def create_async_sandbox_session(
*,
session_type: None = None,
store: AsyncStore | None = None,
client: ContreeAsyncClient | None = None,
client_type: type[ContreeAsyncClient] | None = None,
file_transfer: AsyncFileTransfer | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncSandboxSessionContext[AsyncLazySession]: ...
@overload
def create_async_sandbox_session(
*,
session_type: type[AsyncSessionT],
store: AsyncStore | None = None,
client: ContreeAsyncClient | None = None,
client_type: type[ContreeAsyncClient] | None = None,
file_transfer: AsyncFileTransfer | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncSandboxSessionContext[AsyncSessionT]: ...
[docs]
def create_async_sandbox_session(
*,
session_type: type[AsyncExecutor] | None = None,
session_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncSandboxSessionContext[Any]:
"""Create an async-only context; its default session is AsyncLazySession.
Pass constructor options for the selected class through session_kwargs.
Other keyword arguments configure the factory or its client.
Returns:
An asynchronous context with the configuration accepted by session_maker.
"""
return session_maker(AsyncSandboxSessionContext, session_type=session_type, session_kwargs=session_kwargs, **kwargs)
[docs]
def session_maker( # noqa: PLR0913 - explicit routing of session, profile, and client settings
context_type: Callable[[SandboxConfig], ContextT],
/,
*,
image: str | None = None,
url: str | None = None,
project: str | None = None,
token: str | None = None,
auth_type: Literal["iam", "jwt"] | None = None,
profile: Profile | str | None = None,
config_path: str | Path | None = None,
session_type: type[SyncExecutor | AsyncExecutor] | None = None,
session_id: str | None = None,
node_id: str | None = None,
store: SyncStore | AsyncStore | None = None,
cache_ttl: float = 60,
cwd: str | None = None,
env: dict[str, str] | None = None,
commit_policy: AbstractCommitPolicy | None = None,
file_transfer: SyncFileTransfer | AsyncFileTransfer | None = None,
session_kwargs: Mapping[str, Any] | None = None,
client: Client | None = None,
client_type: type[ContreeSyncClient | ContreeAsyncClient] | None = None,
**client_kwargs: Any,
) -> ContextT:
"""Prepare a sync/async sandbox context without opening connections or a VM.
Without session_type, entry selects LazySession or AsyncLazySession. Explicit
classes must implement SyncExecutor or AsyncExecutor and match the context
mode. Lazy subclasses receive a base session; other classes receive a client
and the standard ContreeSession constructor keywords. session_kwargs adds
arguments to the selected class, never to the underlying lazy base session.
Factory-managed arguments take priority over duplicates in session_kwargs.
Connection arguments, an explicit profile/config_path, and a borrowed client
are separate sources. Without an explicit source, use the client's environment
profile, then its default INI resolution. client_kwargs is forwarded unchanged
to the selected client's from_profile constructor. Created clients use
RetryPolicy() unless retry is supplied; retry=None disables retries.
The default lazy policy snapshots after 10 idle seconds. Normal exit saves
pending work; exceptional exit discards unsaved work. Snapshots retain files,
not running processes. ContreeSession and ContreeAsyncSession save each command
by default in a separate VM. Pass disposable=True to leave their history unchanged.
Returns:
A single-use context owning its created client, memory store, and session.
Supplied clients, stores, and file transfers remain caller-owned.
"""
return context_type(
SandboxConfig(
image=image,
url=url,
project=project,
token=token,
auth_type=auth_type,
profile=profile,
config_path=config_path,
session_type=session_type,
session_id=session_id,
node_id=node_id,
store=store,
cache_ttl=cache_ttl,
cwd=cwd,
env=dict(env) if env is not None else None,
commit_policy=commit_policy,
file_transfer=file_transfer,
session_kwargs=dict(session_kwargs or {}),
client=client,
client_type=client_type,
client_kwargs=client_kwargs,
)
)