blaxel 0.2.36__py3-none-any.whl → 0.2.37__py3-none-any.whl

blaxel/core/sandbox/client/models/upgrade_request.py

@@ -0,0 +1,71 @@
+from typing import Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="UpgradeRequest")
+
+
+@_attrs_define
+class UpgradeRequest:
+    """
+    Attributes:
+        base_url (Union[Unset, str]): Base URL for releases (useful for forks) Example: https://github.com/blaxel-
+            ai/sandbox/releases.
+        version (Union[Unset, str]): Version to upgrade to: "develop", "main", "latest", or specific tag like "v1.0.0"
+            Example: develop.
+    """
+
+    base_url: Union[Unset, str] = UNSET
+    version: Union[Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        base_url = self.base_url
+
+        version = self.version
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if base_url is not UNSET:
+            field_dict["baseUrl"] = base_url
+        if version is not UNSET:
+            field_dict["version"] = version
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T | None:
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        base_url = d.pop("baseUrl", d.pop("base_url", UNSET))
+
+        version = d.pop("version", UNSET)
+
+        upgrade_request = cls(
+            base_url=base_url,
+            version=version,
+        )
+
+        upgrade_request.additional_properties = d
+        return upgrade_request
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

blaxel/core/sandbox/client/models/upgrade_status.py

@@ -0,0 +1,125 @@
+from typing import Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..models.process_upgrade_state import ProcessUpgradeState
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="UpgradeStatus")
+
+
+@_attrs_define
+class UpgradeStatus:
+    """
+    Attributes:
+        status (ProcessUpgradeState):
+        step (str): Current/last step (none, starting, download, validate, replace, completed, skipped) Example:
+            download.
+        version (str): Version being upgraded to Example: latest.
+        binary_path (Union[Unset, str]): Path to downloaded binary Example: /tmp/sandbox-api-new.
+        bytes_downloaded (Union[Unset, int]): Bytes downloaded Example: 25034936.
+        download_url (Union[Unset, str]): URL used for download Example: https://github.com/....
+        error (Union[Unset, str]): Error message if failed Example: Failed to download binary.
+        last_attempt (Union[Unset, str]): When the upgrade was attempted
+    """
+
+    status: ProcessUpgradeState
+    step: str
+    version: str
+    binary_path: Union[Unset, str] = UNSET
+    bytes_downloaded: Union[Unset, int] = UNSET
+    download_url: Union[Unset, str] = UNSET
+    error: Union[Unset, str] = UNSET
+    last_attempt: Union[Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        status = self.status.value
+
+        step = self.step
+
+        version = self.version
+
+        binary_path = self.binary_path
+
+        bytes_downloaded = self.bytes_downloaded
+
+        download_url = self.download_url
+
+        error = self.error
+
+        last_attempt = self.last_attempt
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update(
+            {
+                "status": status,
+                "step": step,
+                "version": version,
+            }
+        )
+        if binary_path is not UNSET:
+            field_dict["binaryPath"] = binary_path
+        if bytes_downloaded is not UNSET:
+            field_dict["bytesDownloaded"] = bytes_downloaded
+        if download_url is not UNSET:
+            field_dict["downloadUrl"] = download_url
+        if error is not UNSET:
+            field_dict["error"] = error
+        if last_attempt is not UNSET:
+            field_dict["lastAttempt"] = last_attempt
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T | None:
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        status = ProcessUpgradeState(d.pop("status"))
+
+        step = d.pop("step")
+
+        version = d.pop("version")
+
+        binary_path = d.pop("binaryPath", d.pop("binary_path", UNSET))
+
+        bytes_downloaded = d.pop("bytesDownloaded", d.pop("bytes_downloaded", UNSET))
+
+        download_url = d.pop("downloadUrl", d.pop("download_url", UNSET))
+
+        error = d.pop("error", UNSET)
+
+        last_attempt = d.pop("lastAttempt", d.pop("last_attempt", UNSET))
+
+        upgrade_status = cls(
+            status=status,
+            step=step,
+            version=version,
+            binary_path=binary_path,
+            bytes_downloaded=bytes_downloaded,
+            download_url=download_url,
+            error=error,
+            last_attempt=last_attempt,
+        )
+
+        upgrade_status.additional_properties = d
+        return upgrade_status
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

blaxel/core/sandbox/default/__init__.py

@@ -7,6 +7,7 @@ from .sandbox import (
     SandboxPreviews,
     SandboxProcess,
 )
+from .system import SandboxSystem
 
 __all__ = [
     "SandboxInstance",
@@ -15,5 +16,6 @@ __all__ = [
     "SandboxPreviews",
     "SandboxProcess",
     "SandboxCodegen",
+    "SandboxSystem",
     "CodeInterpreter",
 ]

blaxel/core/sandbox/default/filesystem.py

@@ -10,6 +10,7 @@ import httpx
 from ...common.settings import settings
 from ..client.models import Directory, FileRequest, SuccessResponse
 from ..types import (
+    AsyncWatchHandle,
     CopyResponse,
     SandboxConfiguration,
     SandboxFilesystemFile,
@@ -327,9 +328,7 @@ class SandboxFileSystem(SandboxAction):
             data = json.loads(await response.aread())
             self.handle_response_error(response)
 
-            from ..client.models.content_search_response import (
-                ContentSearchResponse,
-            )
+            from ..client.models.content_search_response import ContentSearchResponse
 
             return ContentSearchResponse.from_dict(data)
         finally:
@@ -364,8 +363,23 @@ class SandboxFileSystem(SandboxAction):
         path: str,
         callback: Callable[[WatchEvent], None],
         options: Dict[str, Any] | None = None,
-    ) -> Dict[str, Callable]:
-        """Watch for file system changes."""
+    ) -> AsyncWatchHandle:
+        """Watch for file system changes.
+
+        Returns an AsyncWatchHandle that can be used as a context manager:
+
+            async with sandbox.fs.watch(path, callback) as handle:
+                # do something
+            # handle is automatically closed
+
+        Or manually:
+
+            handle = sandbox.fs.watch(path, callback)
+            try:
+                # do something
+            finally:
+                handle.close()
+        """
         path = self.format_path(path)
         closed = False
 
@@ -444,7 +458,7 @@ class SandboxFileSystem(SandboxAction):
             closed = True
             task.cancel()
 
-        return {"close": close}
+        return AsyncWatchHandle(close)
 
     def format_path(self, path: str) -> str:
         """Format path for filesystem operations.

blaxel/core/sandbox/default/preview.py

@@ -1,3 +1,4 @@
+import asyncio
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, Dict, List, Union
@@ -15,6 +16,9 @@ from ...client.api.compute.delete_sandbox_preview import (
 from ...client.api.compute.delete_sandbox_preview_token import (
     asyncio as delete_sandbox_preview_token,
 )
+from ...client.api.compute.get_sandbox_preview import (
+    asyncio_detailed as get_sandbox_preview_detailed,
+)
 from ...client.api.compute.get_sandbox_preview import (
     asyncio as get_sandbox_preview,
 )
@@ -176,15 +180,58 @@ class SandboxPreviews:
         )
         return SandboxPreview(response)
 
-    async def delete(self, preview_name: str) -> dict:
+    async def delete(self, preview_name: str) -> Preview:
         """Delete a preview."""
         response: Preview = await delete_sandbox_preview(
             self.sandbox_name,
             preview_name,
             client=client,
         )
+
+        # If the preview is in DELETING state, wait for it to be fully deleted
+        if response and response.status == "DELETING":
+            await self._wait_for_deletion(preview_name)
+
         return response
 
+    async def _wait_for_deletion(self, preview_name: str, timeout_ms: int = 10000) -> None:
+        """Wait for a preview to be fully deleted.
+
+        Args:
+            preview_name: Name of the preview to wait for
+            timeout_ms: Timeout in milliseconds (default: 10000)
+
+        Raises:
+            Exception: If the preview is still in DELETING state after timeout
+        """
+        print(f"Waiting for preview deletion: {preview_name}")
+        poll_interval = 0.5  # Poll every 500ms
+        elapsed = 0.0
+        timeout_seconds = timeout_ms / 1000.0
+
+        while elapsed < timeout_seconds:
+            try:
+                response = await get_sandbox_preview_detailed(
+                    self.sandbox_name,
+                    preview_name,
+                    client=client,
+                )
+                if response.status_code == 404:
+                    return
+            except errors.UnexpectedStatus as e:
+                # 404 means the preview is deleted
+                if e.status_code == 404:
+                    return
+                raise
+            # Preview still exists, wait and retry
+            await asyncio.sleep(poll_interval)
+            elapsed += poll_interval
+
+        # Timeout reached, but deletion was initiated
+        raise Exception(
+            f"Preview deletion timeout: {preview_name} is still in DELETING state after {timeout_ms}ms"
+        )
+
 
 def to_utc_z(dt: datetime) -> str:
     """Convert datetime to UTC Z format string."""

blaxel/core/sandbox/default/process.py

@@ -6,7 +6,12 @@ import httpx
 from ...common.settings import settings
 from ..client.models import ProcessResponse, SuccessResponse
 from ..client.models.process_request import ProcessRequest
-from ..types import ProcessRequestWithLog, ProcessResponseWithLog, SandboxConfiguration
+from ..types import (
+    AsyncStreamHandle,
+    ProcessRequestWithLog,
+    ProcessResponseWithLog,
+    SandboxConfiguration,
+)
 from .action import SandboxAction
 
 
@@ -18,13 +23,28 @@ class SandboxProcess(SandboxAction):
         self,
         process_name: str,
         options: Dict[str, Callable[[str], None]] | None = None,
-    ) -> Dict[str, Callable[[], None]]:
-        """Stream logs from a process with automatic reconnection and deduplication."""
+    ) -> AsyncStreamHandle:
+        """Stream logs from a process with automatic reconnection and deduplication.
+
+        Returns an AsyncStreamHandle that can be used as a context manager:
+
+            async with sandbox.process.stream_logs(name, options) as handle:
+                # do something
+            # handle is automatically closed
+
+        Or manually:
+
+            handle = sandbox.process.stream_logs(name, options)
+            try:
+                # do something
+            finally:
+                handle.close()
+        """
         if options is None:
             options = {}
 
         reconnect_interval = 30  # 30 seconds in Python (TypeScript uses milliseconds)
-        current_stream = None
+        current_stream: AsyncStreamHandle | None = None
         is_running = True
         reconnect_timer = None
 
@@ -37,7 +57,7 @@ class SandboxProcess(SandboxAction):
 
             # Close existing stream if any
             if current_stream:
-                current_stream["close"]()
+                current_stream.close()
 
             # Create wrapper options with deduplication
             wrapped_options = {}
@@ -104,19 +124,19 @@ class SandboxProcess(SandboxAction):
 
             # Close current stream
             if current_stream:
-                current_stream["close"]()
+                current_stream.close()
                 current_stream = None
 
             # Clear seen logs
             seen_logs.clear()
 
-        return {"close": close}
+        return AsyncStreamHandle(close)
 
     def _stream_logs(
         self,
         identifier: str,
         options: Dict[str, Callable[[str], None]] | None = None,
-    ) -> Dict[str, Callable[[], None]]:
+    ) -> AsyncStreamHandle:
         """Private method to stream logs from a process with callbacks for different output types."""
         if options is None:
             options = {}
@@ -165,7 +185,9 @@ class SandboxProcess(SandboxAction):
                                         options["on_log"](line)
             except Exception as e:
                 # Suppress AbortError when closing
-                if not (hasattr(e, "name") and e.name == "AbortError"):
+                if hasattr(e, "name") and getattr(e, "name") == "AbortError":
+                    pass
+                else:
                     raise e
 
         # Start streaming in the background
@@ -176,7 +198,7 @@ class SandboxProcess(SandboxAction):
             closed = True
             task.cancel()
 
-        return {"close": close}
+        return AsyncStreamHandle(close)
 
     async def exec(
         self,
@@ -203,7 +225,9 @@ class SandboxProcess(SandboxAction):
             if "on_stderr" in process:
                 on_stderr = process["on_stderr"]
                 del process["on_stderr"]
-            process = ProcessRequest.from_dict(process)
+            tmp_process = ProcessRequest.from_dict(process)
+            assert tmp_process is not None
+            process = tmp_process
 
         # Store original wait_for_completion setting
         should_wait_for_completion = process.wait_for_completion
@@ -221,18 +245,27 @@ class SandboxProcess(SandboxAction):
                 self.handle_response_error(response)
                 import json
 
-                response_data = json.loads(content_bytes) if content_bytes else None
-                result = ProcessResponse.from_dict(response_data)
+                if content_bytes:
+                    response_data = json.loads(content_bytes)
+                    result = ProcessResponse.from_dict(response_data)
+                    assert result is not None
+                else:
+                    raise Exception("No content received from response")
             finally:
                 await response.aclose()
 
             if on_log or on_stdout or on_stderr:
-                stream_control = self._stream_logs(
-                    result.pid, {"on_log": on_log, "on_stdout": on_stdout, "on_stderr": on_stderr}
-                )
+                stream_options: dict[str, Callable[[str], None]] = {}
+                if on_log:
+                    stream_options["on_log"] = on_log
+                if on_stdout:
+                    stream_options["on_stdout"] = on_stdout
+                if on_stderr:
+                    stream_options["on_stderr"] = on_stderr
+                stream_control = self._stream_logs(result.pid, stream_options)
                 return ProcessResponseWithLog(
                     result,
-                    lambda: stream_control["close"]() if stream_control else None,
+                    lambda: stream_control.close() if stream_control else None,
                 )
 
             return result
@@ -277,6 +310,7 @@ class SandboxProcess(SandboxAction):
                     content = await response.aread()
                     data = json.loads(content)
                     result = ProcessResponse.from_dict(data)
+                    assert result is not None
 
                     # If process already completed (server waited), emit logs through callbacks
                     if result.status == "completed" or result.status == "failed":
@@ -378,7 +412,9 @@ class SandboxProcess(SandboxAction):
         try:
             data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return ProcessResponse.from_dict(data)
+            result = ProcessResponse.from_dict(data)
+            assert result is not None
+            return result
         finally:
             await response.aclose()
 
@@ -390,7 +426,12 @@ class SandboxProcess(SandboxAction):
         try:
             data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return [ProcessResponse.from_dict(item) for item in data]
+            results = []
+            for item in data:
+                result = ProcessResponse.from_dict(item)
+                assert result is not None
+                results.append(result)
+            return results
         finally:
             await response.aclose()
 
@@ -402,7 +443,9 @@ class SandboxProcess(SandboxAction):
         try:
             data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return SuccessResponse.from_dict(data)
+            result = SuccessResponse.from_dict(data)
+            assert result is not None
+            return result
         finally:
             await response.aclose()
 
@@ -414,7 +457,9 @@ class SandboxProcess(SandboxAction):
         try:
             data = json.loads(await response.aread())
             self.handle_response_error(response)
-            return SuccessResponse.from_dict(data)
+            result = SuccessResponse.from_dict(data)
+            assert result is not None
+            return result
         finally:
             await response.aclose()
 

blaxel/core/sandbox/default/sandbox.py

@@ -8,7 +8,14 @@ from ...client.api.compute.get_sandbox import asyncio as get_sandbox
 from ...client.api.compute.list_sandboxes import asyncio as list_sandboxes
 from ...client.api.compute.update_sandbox import asyncio as update_sandbox
 from ...client.client import client
-from ...client.models import Metadata, Sandbox, SandboxLifecycle, SandboxRuntime, SandboxSpec
+from ...client.models import (
+    Metadata,
+    MetadataLabels,
+    Sandbox,
+    SandboxLifecycle,
+    SandboxRuntime,
+    SandboxSpec,
+)
 from ...client.models.error import Error
 from ...client.models.sandbox_error import SandboxError
 from ...client.types import UNSET
@@ -25,6 +32,7 @@ from .network import SandboxNetwork
 from .preview import SandboxPreviews
 from .process import SandboxProcess
 from .session import SandboxSessions
+from .system import SandboxSystem
 
 
 class SandboxAPIError(Exception):
@@ -87,6 +95,7 @@ class SandboxInstance:
         self.sessions = SandboxSessions(self.config)
         self.network = SandboxNetwork(self.config)
         self.codegen = SandboxCodegen(self.config)
+        self.system = SandboxSystem(self.config)
 
     @property
     def metadata(self):
@@ -104,6 +113,14 @@ class SandboxInstance:
     def spec(self):
         return self.sandbox.spec
 
+    @property
+    def last_used_at(self):
+        return self.sandbox.last_used_at
+
+    @property
+    def expires_in(self):
+        return self.sandbox.expires_in
+
     async def wait(self, max_wait: int = 60000, interval: int = 1000) -> "SandboxInstance":
         logger.warning(
             "⚠️  Warning: sandbox.wait() is deprecated. You don't need to wait for the sandbox to be deployed anymore."
@@ -167,9 +184,13 @@ class SandboxInstance:
             lifecycle = config.lifecycle
             # snapshot_enabled = sandbox.snapshot_enabled
 
+            labels = MetadataLabels.from_dict(config.labels) if config.labels else UNSET
+            if labels is None:
+                labels = UNSET
+
             # Create full Sandbox object
             sandbox = Sandbox(
-                metadata=Metadata(name=name, labels=config.labels),
+                metadata=Metadata(name=name, labels=labels),
                 spec=SandboxSpec(
                     runtime=SandboxRuntime(
                         image=image,
@@ -182,19 +203,28 @@ class SandboxInstance:
             )
 
             # Set ttl and expires if provided
-            if ttl:
+            if ttl and sandbox.spec.runtime:
                 sandbox.spec.runtime.ttl = ttl
-            if expires:
+            if expires and sandbox.spec.runtime:
                 sandbox.spec.runtime.expires = expires.isoformat()
             if region:
                 sandbox.spec.region = region
             if lifecycle:
-                sandbox.spec.lifecycle = lifecycle
+                if type(lifecycle) is dict:
+                    lifecycle = SandboxLifecycle.from_dict(lifecycle)
+                    assert lifecycle is not None
+                    sandbox.spec.lifecycle = lifecycle
+                elif type(lifecycle) is SandboxLifecycle:
+                    sandbox.spec.lifecycle = lifecycle
+                else:
+                    raise ValueError(f"Invalid lifecycle type: {type(lifecycle)}")
         else:
             # Handle existing Sandbox object or dict conversion
             if isinstance(sandbox, dict):
                 sandbox = Sandbox.from_dict(sandbox)
+                assert sandbox is not None
 
+            assert isinstance(sandbox, Sandbox)
             # Set defaults for missing fields
             if not sandbox.metadata:
                 sandbox.metadata = Metadata(name=default_name)
@@ -220,6 +250,7 @@ class SandboxInstance:
             message = response.message if response.message else str(response)
             raise SandboxAPIError(message, status_code=status_code, code=code)
 
+        assert response is not None
         instance = cls(response)
         # TODO remove this part once we have a better way to handle this
         if safe: