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

blaxel/core/sandbox/sync/process.py

@@ -7,7 +7,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 (
+    ProcessRequestWithLog,
+    ProcessResponseWithLog,
+    SandboxConfiguration,
+    StreamHandle,
+)
 from .action import SyncSandboxAction
 
 
@@ -19,19 +24,35 @@ class SyncSandboxProcess(SyncSandboxAction):
         self,
         process_name: str,
         options: Dict[str, Callable[[str], None]] | None = None,
-    ) -> Dict[str, Callable[[], None]]:
+    ) -> StreamHandle:
+        """Stream logs from a process with automatic reconnection and deduplication.
+
+        Returns a StreamHandle that can be used as a context manager:
+
+            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
         is_running = threading.Event()
         is_running.set()
         seen_logs = set()
-        current_close = {"fn": None}
+        current_stream: StreamHandle | None = None
         timer_lock = threading.Lock()
-        reconnect_timer = {"t": None}
+        reconnect_timer: dict[str, threading.Timer | None] = {"t": None}
 
         def start_stream():
-            nonlocal current_close
+            nonlocal current_stream
             log_counter = [0]
 
             def make_dedup(cb_key: str):
@@ -52,9 +73,9 @@ class SyncSandboxProcess(SyncSandboxAction):
                 wrapped_options["on_stdout"] = make_dedup("on_stdout")
             if "on_stderr" in options:
                 wrapped_options["on_stderr"] = make_dedup("on_stderr")
-            if current_close["fn"]:
-                current_close["fn"]()
-            current_close["fn"] = self._stream_logs(process_name, wrapped_options)["close"]
+            if current_stream:
+                current_stream.close()
+            current_stream = self._stream_logs(process_name, wrapped_options)
 
             def schedule():
                 if is_running.is_set():
@@ -71,23 +92,25 @@ class SyncSandboxProcess(SyncSandboxAction):
         start_stream()
 
         def close():
+            nonlocal current_stream
             is_running.clear()
             with timer_lock:
                 if reconnect_timer["t"]:
                     reconnect_timer["t"].cancel()
                     reconnect_timer["t"] = None
-            if current_close["fn"]:
-                current_close["fn"]()
-                current_close["fn"] = None
+            if current_stream:
+                current_stream.close()
+                current_stream = None
             seen_logs.clear()
 
-        return {"close": close}
+        return StreamHandle(close)
 
     def _stream_logs(
         self,
         identifier: str,
         options: Dict[str, Callable[[str], None]] | None = None,
-    ) -> Dict[str, Callable[[], None]]:
+    ) -> StreamHandle:
+        """Private method to stream logs from a process with callbacks for different output types."""
         if options is None:
             options = {}
         closed = threading.Event()
@@ -136,7 +159,7 @@ class SyncSandboxProcess(SyncSandboxAction):
         def close():
             closed.set()
 
-        return {"close": close}
+        return StreamHandle(close)
 
     def exec(
         self,
@@ -191,7 +214,7 @@ class SyncSandboxProcess(SyncSandboxAction):
                     )
                     return ProcessResponseWithLog(
                         result,
-                        lambda: stream_control["close"]() if stream_control else None,
+                        lambda: stream_control.close() if stream_control else None,
                     )
 
                 return result

blaxel/core/sandbox/sync/sandbox.py

@@ -8,10 +8,17 @@ from ...client.api.compute.get_sandbox import sync as get_sandbox
 from ...client.api.compute.list_sandboxes import sync as list_sandboxes
 from ...client.api.compute.update_sandbox import sync as update_sandbox
 from ...client.client import client
-from ...client.models import Metadata, Sandbox, SandboxRuntime, SandboxSpec
+from ...client.models import (
+    Metadata,
+    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 ..default.sandbox import SandboxAPIError
 from ..types import (
     SandboxConfiguration,
@@ -25,6 +32,7 @@ from .network import SyncSandboxNetwork
 from .preview import SyncSandboxPreviews
 from .process import SyncSandboxProcess
 from .session import SyncSandboxSessions
+from .system import SyncSandboxSystem
 
 logger = logging.getLogger(__name__)
 
@@ -72,6 +80,7 @@ class SyncSandboxInstance:
         self.sessions = SyncSandboxSessions(self.config)
         self.network = SyncSandboxNetwork(self.config)
         self.codegen = SyncSandboxCodegen(self.config)
+        self.system = SyncSandboxSystem(self.config)
 
     @property
     def metadata(self):
@@ -89,6 +98,14 @@ class SyncSandboxInstance:
     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
+
     def wait(self, max_wait: int = 60000, interval: int = 1000) -> "SyncSandboxInstance":
         logger.warning(
             "⚠️  Warning: sandbox.wait() is deprecated. You don't need to wait for the sandbox to be deployed anymore."
@@ -144,7 +161,7 @@ class SyncSandboxInstance:
             volumes = config._normalize_volumes() or UNSET
             ttl = config.ttl
             expires = config.expires
-            region = config.region
+            region = config.region or settings.region
             lifecycle = config.lifecycle
             sandbox = Sandbox(
                 metadata=Metadata(name=name, labels=config.labels),
@@ -158,17 +175,26 @@ class SyncSandboxInstance:
                     volumes=volumes,
                 ),
             )
-            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:
             if isinstance(sandbox, dict):
                 sandbox = Sandbox.from_dict(sandbox)
+            assert isinstance(sandbox, Sandbox)
+
             if not sandbox.metadata:
                 sandbox.metadata = Metadata(name=default_name)
             if not sandbox.spec:
@@ -250,6 +276,72 @@ class SyncSandboxInstance:
         )
         return cls(response)
 
+    @classmethod
+    def update_ttl(cls, sandbox_name: str, ttl: str) -> "SyncSandboxInstance":
+        """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 SyncSandboxInstance with updated TTL
+        """
+        # Get the existing sandbox
+        sandbox_instance = 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 = update_sandbox(
+            sandbox_name=sandbox_name,
+            client=client,
+            body=updated_sandbox,
+        )
+
+        return cls(response)
+
+    @classmethod
+    def update_lifecycle(
+        cls, sandbox_name: str, lifecycle: SandboxLifecycle
+    ) -> "SyncSandboxInstance":
+        """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 SyncSandboxInstance with updated lifecycle
+        """
+        # Get the existing sandbox
+        sandbox_instance = 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 = update_sandbox(
+            sandbox_name=sandbox_name,
+            client=client,
+            body=updated_sandbox,
+        )
+
+        return cls(response)
+
     @classmethod
     def create_if_not_exists(
         cls, sandbox: Union[Sandbox, SandboxCreateConfiguration, Dict[str, Any]]

blaxel/core/sandbox/sync/system.py

@@ -0,0 +1,71 @@
+from typing import Union
+
+from ...common.settings import settings
+from ..client.api.system.get_health import sync as get_health
+from ..client.api.system.post_upgrade import sync as post_upgrade
+from ..client.client import Client
+from ..client.models import ErrorResponse, HealthResponse, SuccessResponse, UpgradeRequest
+from ..types import SandboxConfiguration
+from .action import SyncSandboxAction
+
+
+class SyncSandboxSystem(SyncSandboxAction):
+    """System operations for sandbox including upgrade functionality (sync version)."""
+
+    def __init__(self, sandbox_config: SandboxConfiguration):
+        super().__init__(sandbox_config)
+
+    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},
+        )
+
+        with client:
+            response = 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
+
+    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},
+        )
+
+        with client:
+            response = get_health(client=client)
+            if response is None:
+                raise Exception("Failed to get health status")
+            return response

blaxel/core/sandbox/types.py

@@ -4,7 +4,14 @@ from typing import Any, Callable, Dict, List, TypeVar, Union
 import httpx
 from attrs import define as _attrs_define
 
-from ..client.models import Port, Sandbox, SandboxLifecycle, VolumeAttachment
+from ..client.models import (
+    Env,
+    Port,
+    PortProtocol,
+    Sandbox,
+    SandboxLifecycle,
+    VolumeAttachment,
+)
 from ..client.types import UNSET
 from .client.models.process_request import ProcessRequest
 from .client.models.process_response import ProcessResponse
@@ -205,7 +212,7 @@ class SandboxCreateConfiguration:
             if isinstance(port, Port):
                 # If it's already a Port object, ensure protocol defaults to HTTP
                 if port.protocol is UNSET or not port.protocol:
-                    port.protocol = "HTTP"
+                    port.protocol = PortProtocol.HTTP
                 port_objects.append(port)
             elif isinstance(port, dict):
                 # Convert dict to Port object with HTTP as default protocol
@@ -218,20 +225,22 @@ class SandboxCreateConfiguration:
 
         return port_objects
 
-    def _normalize_envs(self) -> List[Dict[str, str]] | None:
+    def _normalize_envs(self) -> List[Env] | None:
         """Convert envs to list of dicts with name and value keys."""
         if not self.envs:
             return None
 
         env_objects = []
         for env in self.envs:
-            if isinstance(env, dict):
+            if isinstance(env, Env):
+                env_objects.append(env)
+            elif isinstance(env, dict):
                 # Validate that the dict has the required keys
                 if "name" not in env or "value" not in env:
                     raise ValueError(
                         f"Environment variable dict must have 'name' and 'value' keys: {env}"
                     )
-                env_objects.append({"name": env["name"], "value": env["value"]})
+                env_objects.append(Env(name=env["name"], value=env["value"]))
             else:
                 raise ValueError(
                     f"Invalid env type: {type(env)}. Expected dict with 'name' and 'value' keys."
@@ -385,3 +394,201 @@ class Context:
     @classmethod
     def from_json(cls, data: Dict[str, Any]) -> "Context":
         return cls(id=str(data.get("id") or data.get("context_id") or ""))
+
+
+class StreamHandle:
+    """Handle for managing a streaming operation (sync version).
+
+    Can be used as a context manager for automatic cleanup:
+
+        with sandbox.process.stream_logs(name, options) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.process.stream_logs(name, options)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the stream and stop receiving data."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the stream handle has been closed."""
+        return self._closed
+
+    def __enter__(self) -> "StreamHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)
+
+
+class AsyncStreamHandle:
+    """Handle for managing a streaming operation (async version).
+
+    Can be used as an async context manager for automatic cleanup:
+
+        async with sandbox.process.stream_logs(name, options) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.process.stream_logs(name, options)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the stream and stop receiving data."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the stream handle has been closed."""
+        return self._closed
+
+    async def __aenter__(self) -> "AsyncStreamHandle":
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Also support sync context manager for convenience
+    def __enter__(self) -> "AsyncStreamHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)
+
+
+class WatchHandle:
+    """Handle for managing a file system watch operation (sync version).
+
+    Can be used as a context manager for automatic cleanup:
+
+        with sandbox.fs.watch(path, callback) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.fs.watch(path, callback)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the watch and stop receiving events."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the watch handle has been closed."""
+        return self._closed
+
+    def __enter__(self) -> "WatchHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)
+
+
+class AsyncWatchHandle:
+    """Handle for managing a file system watch operation (async version).
+
+    Can be used as an async context manager for automatic cleanup:
+
+        async with sandbox.fs.watch(path, callback) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.fs.watch(path, callback)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the watch and stop receiving events."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the watch handle has been closed."""
+        return self._closed
+
+    async def __aenter__(self) -> "AsyncWatchHandle":
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Also support sync context manager for convenience
+    def __enter__(self) -> "AsyncWatchHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)