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

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,10 +8,18 @@ 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, 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
+from ...common.settings import settings
 from ..types import (
     SandboxConfiguration,
     SandboxCreateConfiguration,
@@ -24,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):
@@ -86,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):
@@ -103,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."
@@ -162,13 +180,17 @@ class SandboxInstance:
             volumes = config._normalize_volumes() or UNSET
             ttl = config.ttl
             expires = config.expires
-            region = config.region
+            region = config.region or settings.region
             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,
@@ -181,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)
@@ -219,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:
@@ -304,6 +336,72 @@ class SandboxInstance:
         # Return new instance with updated sandbox
         return cls(response)
 
+    @classmethod
+    async def update_ttl(cls, sandbox_name: str, ttl: str) -> "SandboxInstance":
+        """Update sandbox TTL without recreating it.
+
+        Args:
+            sandbox_name: The name of the sandbox to update
+            ttl: The new TTL value (e.g., "5m", "1h", "30s")
+
+        Returns:
+            A new SandboxInstance with updated TTL
+        """
+        # Get the existing sandbox
+        sandbox_instance = await cls.get(sandbox_name)
+        sandbox = sandbox_instance.sandbox
+
+        # Prepare the updated sandbox object
+        updated_sandbox = Sandbox.from_dict(sandbox.to_dict())
+        if updated_sandbox.spec is None or updated_sandbox.spec.runtime is None:
+            raise ValueError(f"Sandbox {sandbox_name} has invalid spec")
+
+        # Update TTL
+        updated_sandbox.spec.runtime.ttl = ttl
+
+        # Call the update API
+        response = await update_sandbox(
+            sandbox_name=sandbox_name,
+            client=client,
+            body=updated_sandbox,
+        )
+
+        return cls(response)
+
+    @classmethod
+    async def update_lifecycle(
+        cls, sandbox_name: str, lifecycle: SandboxLifecycle
+    ) -> "SandboxInstance":
+        """Update sandbox lifecycle configuration without recreating it.
+
+        Args:
+            sandbox_name: The name of the sandbox to update
+            lifecycle: The new lifecycle configuration
+
+        Returns:
+            A new SandboxInstance with updated lifecycle
+        """
+        # Get the existing sandbox
+        sandbox_instance = await cls.get(sandbox_name)
+        sandbox = sandbox_instance.sandbox
+
+        # Prepare the updated sandbox object
+        updated_sandbox = Sandbox.from_dict(sandbox.to_dict())
+        if updated_sandbox.spec is None:
+            raise ValueError(f"Sandbox {sandbox_name} has invalid spec")
+
+        # Update lifecycle
+        updated_sandbox.spec.lifecycle = lifecycle
+
+        # Call the update API
+        response = await update_sandbox(
+            sandbox_name=sandbox_name,
+            client=client,
+            body=updated_sandbox,
+        )
+
+        return cls(response)
+
     @classmethod
     async def create_if_not_exists(
         cls, sandbox: Union[Sandbox, SandboxCreateConfiguration, Dict[str, Any]]

blaxel/core/sandbox/default/system.py

@@ -0,0 +1,71 @@
+from typing import Union
+
+from ...common.settings import settings
+from ..client.api.system.get_health import asyncio as get_health
+from ..client.api.system.post_upgrade import asyncio as post_upgrade
+from ..client.client import Client
+from ..client.models import ErrorResponse, HealthResponse, SuccessResponse, UpgradeRequest
+from ..types import SandboxConfiguration
+from .action import SandboxAction
+
+
+class SandboxSystem(SandboxAction):
+    """System operations for sandbox including upgrade functionality."""
+
+    def __init__(self, sandbox_config: SandboxConfiguration):
+        super().__init__(sandbox_config)
+
+    async def upgrade(
+        self,
+        version: str | None = None,
+        base_url: str | None = None,
+    ) -> SuccessResponse:
+        """Upgrade the sandbox-api to a new version.
+
+        Triggers an upgrade of the sandbox-api process. Returns immediately before upgrading.
+        The upgrade will: download the specified binary from GitHub releases, validate it, and restart.
+        All running processes will be preserved across the upgrade.
+
+        Args:
+            version: Version to upgrade to - "develop" (default), "main", "latest",
+                     or specific tag like "v1.0.0"
+            base_url: Base URL for releases (useful for forks, defaults to
+                      https://github.com/blaxel-ai/sandbox/releases)
+
+        Returns:
+            SuccessResponse with status information
+        """
+        request = UpgradeRequest(version=version, base_url=base_url)
+
+        client = Client(
+            base_url=self.url,
+            headers={**settings.headers, **self.sandbox_config.headers},
+        )
+
+        async with client:
+            response = await post_upgrade(client=client, body=request)
+            if response is None:
+                raise Exception("Failed to upgrade sandbox")
+            if isinstance(response, ErrorResponse):
+                raise Exception(f"Upgrade failed: {response.error}")
+            return response
+
+    async def health(self) -> HealthResponse:
+        """Get health status and system information.
+
+        Returns health status and system information including upgrade count and binary details.
+        Also includes last upgrade attempt status with detailed error information if available.
+
+        Returns:
+            HealthResponse with system status information
+        """
+        client = Client(
+            base_url=self.url,
+            headers={**settings.headers, **self.sandbox_config.headers},
+        )
+
+        async with client:
+            response = await get_health(client=client)
+            if response is None:
+                raise Exception("Failed to get health status")
+            return response

blaxel/core/sandbox/sync/__init__.py

@@ -6,6 +6,7 @@ from .sandbox import (
     SyncSandboxPreviews,
     SyncSandboxProcess,
 )
+from .system import SyncSandboxSystem
 
 __all__ = [
     "SyncSandboxInstance",
@@ -13,5 +14,6 @@ __all__ = [
     "SyncSandboxPreviews",
     "SyncSandboxProcess",
     "SyncSandboxCodegen",
+    "SyncSandboxSystem",
     "SyncCodeInterpreter",
 ]

blaxel/core/sandbox/sync/filesystem.py

@@ -14,6 +14,7 @@ from ..types import (
     SandboxConfiguration,
     SandboxFilesystemFile,
     WatchEvent,
+    WatchHandle,
 )
 from .action import SyncSandboxAction
 
@@ -162,7 +163,23 @@ class SyncSandboxFileSystem(SyncSandboxAction):
         path: str,
         callback: Callable[[WatchEvent], None],
         options: Dict[str, Any] | None = None,
-    ) -> Dict[str, Callable]:
+    ) -> WatchHandle:
+        """Watch for file system changes.
+
+        Returns a WatchHandle that can be used as a context manager:
+
+            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 = threading.Event()
         if options is None:
@@ -226,7 +243,7 @@ class SyncSandboxFileSystem(SyncSandboxAction):
         def close():
             closed.set()
 
-        return {"close": close}
+        return WatchHandle(close)
 
     def format_path(self, path: str) -> str:
         return path

blaxel/core/sandbox/sync/preview.py

@@ -1,3 +1,4 @@
+import time
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, Dict, List, Union
@@ -11,6 +12,9 @@ from ...client.api.compute.delete_sandbox_preview import sync as delete_sandbox_
 from ...client.api.compute.delete_sandbox_preview_token import (
     sync as delete_sandbox_preview_token,
 )
+from ...client.api.compute.get_sandbox_preview import (
+    sync_detailed as get_sandbox_preview_detailed,
+)
 from ...client.api.compute.get_sandbox_preview import sync as get_sandbox_preview
 from ...client.api.compute.list_sandbox_preview_tokens import (
     sync as list_sandbox_preview_tokens,
@@ -188,9 +192,52 @@ class SyncSandboxPreviews:
             preview_name,
             client=client,
         )
-        if response:
-            return response
-        raise errors.UnexpectedStatus(400, b"Failed to delete preview")
+        if not response:
+            raise errors.UnexpectedStatus(400, b"Failed to delete preview")
+
+        # If the preview is in DELETING state, wait for it to be fully deleted
+        if response.status == "DELETING":
+            self._wait_for_deletion(preview_name)
+
+        return response
+
+    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 = 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
+            time.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: