blaxel 0.1.22rc70__py3-none-any.whl → 0.2.0rc2__py3-none-any.whl

blaxel/{sandbox → core/sandbox}/preview.py

@@ -1,25 +1,38 @@
 from dataclasses import dataclass
 from datetime import datetime
-from typing import List, Optional
+from typing import Any, Dict, List, Optional, Union
 
-from ..client.api.compute.create_sandbox_preview import asyncio as create_sandbox_preview
+from ..client.api.compute.create_sandbox_preview import (
+    asyncio as create_sandbox_preview,
+)
 from ..client.api.compute.create_sandbox_preview_token import (
     asyncio as create_sandbox_preview_token,
 )
-from ..client.api.compute.delete_sandbox_preview import asyncio as delete_sandbox_preview
+from ..client.api.compute.delete_sandbox_preview import (
+    asyncio as delete_sandbox_preview,
+)
 from ..client.api.compute.delete_sandbox_preview_token import (
     asyncio as delete_sandbox_preview_token,
 )
 from ..client.api.compute.get_sandbox_preview import asyncio as get_sandbox_preview
-from ..client.api.compute.list_sandbox_preview_tokens import asyncio as list_sandbox_preview_tokens
+from ..client.api.compute.list_sandbox_preview_tokens import (
+    asyncio as list_sandbox_preview_tokens,
+)
 from ..client.api.compute.list_sandbox_previews import asyncio as list_sandbox_previews
 from ..client.client import client
-from ..client.models import Preview, PreviewSpec, PreviewToken, PreviewTokenSpec, Sandbox
+from ..client.models import (
+    Preview,
+    PreviewSpec,
+    PreviewToken,
+    PreviewTokenSpec,
+    Sandbox,
+)
 
 
 @dataclass
 class SandboxPreviewToken:
     """Represents a preview token with its value and expiration."""
+
     preview_token: PreviewToken
 
     @property
@@ -30,23 +43,22 @@ class SandboxPreviewToken:
     def expires_at(self) -> datetime:
         return self.preview_token.spec.expires_at if self.preview_token.spec else datetime.now()
 
+
 class SandboxPreviewTokens:
     """Manages preview tokens for a sandbox preview."""
-    def __init__(self, preview: Preview):
+
+    def __init__(self, preview: Preview, sandbox_name: str):
         self.preview = preview
+        self.sandbox_name = sandbox_name
 
     @property
     def preview_name(self) -> str:
         return self.preview.metadata.name if self.preview.metadata else ""
 
-    @property
-    def resource_name(self) -> str:
-        return self.preview.metadata.resource_name if self.preview.metadata else ""
-
     async def create(self, expires_at: datetime) -> SandboxPreviewToken:
         """Create a new preview token."""
         response: PreviewToken = await create_sandbox_preview_token(
-            self.resource_name,
+            self.sandbox_name,
             self.preview_name,
             body=PreviewToken(
                 spec=PreviewTokenSpec(
@@ -60,7 +72,7 @@ class SandboxPreviewTokens:
     async def list(self) -> List[SandboxPreviewToken]:
         """List all preview tokens."""
         response: List[PreviewToken] = await list_sandbox_preview_tokens(
-            self.resource_name,
+            self.sandbox_name,
             self.preview_name,
             client=client,
         )
@@ -69,18 +81,21 @@ class SandboxPreviewTokens:
     async def delete(self, token_name: str) -> dict:
         """Delete a preview token."""
         response: PreviewToken = await delete_sandbox_preview_token(
-            self.resource_name,
+            self.sandbox_name,
             self.preview_name,
             token_name,
             client=client,
         )
         return response
 
+
 class SandboxPreview:
     """Represents a sandbox preview with its metadata and tokens."""
-    def __init__(self, preview: Preview):
+
+    def __init__(self, preview: Preview, sandbox_name: str = ""):
         self.preview = preview
-        self.tokens = SandboxPreviewTokens(self)
+        self.sandbox_name = sandbox_name
+        self.tokens = SandboxPreviewTokens(preview, sandbox_name)
 
     @property
     def name(self) -> str:
@@ -94,8 +109,10 @@ class SandboxPreview:
     def spec(self) -> Optional[PreviewSpec]:
         return self.preview.spec
 
+
 class SandboxPreviews:
     """Manages sandbox previews."""
+
     def __init__(self, sandbox: Sandbox):
         self.sandbox = sandbox
 
@@ -111,8 +128,11 @@ class SandboxPreviews:
         )
         return [SandboxPreview(preview) for preview in response]
 
-    async def create(self, preview: Preview) -> SandboxPreview:
+    async def create(self, preview: Union[Preview, Dict[str, Any]]) -> SandboxPreview:
         """Create a new preview."""
+        if isinstance(preview, dict):
+            preview = Preview.from_dict(preview)
+
         response: Preview = await create_sandbox_preview(
             self.sandbox_name,
             body=preview,
@@ -138,5 +158,13 @@ class SandboxPreviews:
         )
         return response
 
+
 def to_utc_z(dt: datetime) -> str:
-    return dt.isoformat(timespec="milliseconds").replace("+00:00", "Z")
+    """Convert datetime to UTC Z format string."""
+    # Simple approach: format as ISO string and add Z
+    iso_string = dt.isoformat()
+    if iso_string.endswith("+00:00"):
+        return iso_string.replace("+00:00", "Z")
+    elif "T" in iso_string and not iso_string.endswith("Z"):
+        return iso_string + "Z"
+    return iso_string

blaxel/core/sandbox/process.py

@@ -0,0 +1,159 @@
+import asyncio
+from typing import Any, Callable, Dict, Literal, Optional, Union
+
+import httpx
+
+from ..common.settings import settings
+from .action import SandboxAction
+from .client.models import ProcessResponse, SuccessResponse
+from .client.models.process_request import ProcessRequest
+from .types import SandboxConfiguration
+
+
+class SandboxProcess(SandboxAction):
+    def __init__(self, sandbox_config: SandboxConfiguration):
+        super().__init__(sandbox_config)
+
+    def stream_logs(
+        self, identifier: str, options: Optional[Dict[str, Callable[[str], None]]] = None
+    ) -> Dict[str, Callable[[], None]]:
+        """Stream logs from a process with callbacks for different output types."""
+        if options is None:
+            options = {}
+
+        closed = False
+
+        async def start_streaming():
+            nonlocal closed
+
+            url = f"{self.url}/process/{identifier}/logs/stream"
+            headers = {**settings.headers, **self.sandbox_config.headers}
+
+            try:
+                async with httpx.AsyncClient() as client_instance:
+                    async with client_instance.stream("GET", url, headers=headers) as response:
+                        if response.status_code != 200:
+                            raise Exception(f"Failed to stream logs: {await response.aread()}")
+
+                        buffer = ""
+                        async for chunk in response.aiter_text():
+                            if closed:
+                                break
+
+                            buffer += chunk
+                            lines = buffer.split("\n")
+                            buffer = lines.pop()  # Keep incomplete line in buffer
+
+                            for line in lines:
+                                if line.startswith("stdout:"):
+                                    content = line[7:]  # Remove 'stdout:' prefix
+                                    if options.get("on_stdout"):
+                                        options["on_stdout"](content)
+                                    if options.get("on_log"):
+                                        options["on_log"](content)
+                                elif line.startswith("stderr:"):
+                                    content = line[7:]  # Remove 'stderr:' prefix
+                                    if options.get("on_stderr"):
+                                        options["on_stderr"](content)
+                                    if options.get("on_log"):
+                                        options["on_log"](content)
+                                else:
+                                    if options.get("on_log"):
+                                        options["on_log"](line)
+            except Exception as e:
+                # Suppress AbortError when closing
+                if not (hasattr(e, "name") and e.name == "AbortError"):
+                    raise e
+
+        # Start streaming in the background
+        task = asyncio.create_task(start_streaming())
+
+        def close():
+            nonlocal closed
+            closed = True
+            task.cancel()
+
+        return {"close": close}
+
+    async def exec(self, process: Union[ProcessRequest, Dict[str, Any]]) -> ProcessResponse:
+        if isinstance(process, dict):
+            process = ProcessRequest.from_dict(process)
+
+        async with self.get_client() as client_instance:
+            response = await client_instance.post("/process", json=process.to_dict())
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return ProcessResponse.from_dict(response.json())
+
+    async def wait(
+        self, identifier: str, max_wait: int = 60000, interval: int = 1000
+    ) -> ProcessResponse:
+        """Wait for a process to complete."""
+        start_time = asyncio.get_event_loop().time() * 1000  # Convert to milliseconds
+        status = "running"
+        data = await self.get(identifier)
+
+        while status == "running":
+            await asyncio.sleep(interval / 1000)  # Convert to seconds
+            try:
+                data = await self.get(identifier)
+                status = data.status or "running"
+            except:
+                break
+
+            if (asyncio.get_event_loop().time() * 1000) - start_time > max_wait:
+                raise Exception("Process did not finish in time")
+
+        return data
+
+    async def get(self, identifier: str) -> ProcessResponse:
+        async with self.get_client() as client_instance:
+            response = await client_instance.get(f"/process/{identifier}")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return ProcessResponse.from_dict(response.json())
+
+    async def list(self) -> list[ProcessResponse]:
+        async with self.get_client() as client_instance:
+            response = await client_instance.get("/process")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return [ProcessResponse.from_dict(item) for item in response.json()]
+
+    async def stop(self, identifier: str) -> SuccessResponse:
+        async with self.get_client() as client_instance:
+            response = await client_instance.delete(f"/process/{identifier}")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return SuccessResponse.from_dict(response.json())
+
+    async def kill(self, identifier: str) -> SuccessResponse:
+        async with self.get_client() as client_instance:
+            response = await client_instance.delete(f"/process/{identifier}/kill")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+            return SuccessResponse.from_dict(response.json())
+
+    async def logs(
+        self, identifier: str, log_type: Literal["stdout", "stderr", "all"] = "all"
+    ) -> str:
+        async with self.get_client() as client_instance:
+            response = await client_instance.get(f"/process/{identifier}/logs")
+            self.handle_response_error(
+                response, response.json() if response.content else None, None
+            )
+
+            data = response.json()
+            if log_type == "all":
+                return data.get("logs", "")
+            elif log_type == "stdout":
+                return data.get("stdout", "")
+            elif log_type == "stderr":
+                return data.get("stderr", "")
+
+            raise Exception("Unsupported log type")

blaxel/{sandbox → core/sandbox}/sandbox.py

@@ -1,26 +1,33 @@
 import asyncio
 import logging
 import time
-from typing import List
+from typing import Any, Dict, List, Union
 
 from ..client.api.compute.create_sandbox import asyncio as create_sandbox
 from ..client.api.compute.delete_sandbox import asyncio as delete_sandbox
 from ..client.api.compute.get_sandbox import asyncio as get_sandbox
 from ..client.api.compute.list_sandboxes import asyncio as list_sandboxes
 from ..client.client import client
-from ..client.models import Sandbox
+from ..client.models import Metadata, Sandbox
 from .filesystem import SandboxFileSystem
+from .network import SandboxNetwork
 from .preview import SandboxPreviews
 from .process import SandboxProcess
+from .session import SandboxSessions
+from .types import SandboxConfiguration, SessionWithToken
 
 logger = logging.getLogger(__name__)
 
+
 class SandboxInstance:
     def __init__(self, sandbox: Sandbox):
         self.sandbox = sandbox
-        self.fs = SandboxFileSystem(sandbox)
-        self.process = SandboxProcess(sandbox)
+        self.config = SandboxConfiguration(sandbox)
+        self.fs = SandboxFileSystem(self.config)
+        self.process = SandboxProcess(self.config)
         self.previews = SandboxPreviews(sandbox)
+        self.sessions = SandboxSessions(self.config)
+        self.network = SandboxNetwork(self.config)
 
     @property
     def metadata(self):
@@ -49,6 +56,7 @@ class SandboxInstance:
                 )
                 logger.info(f"Waiting for sandbox to be deployed, status: {response.status}")
                 self.sandbox = response
+                self.config = SandboxConfiguration(self.sandbox)
             except Exception as e:
                 logger.error("Could not retrieve sandbox", exc_info=e)
 
@@ -59,7 +67,10 @@ class SandboxInstance:
                 raise Exception("Sandbox did not deploy in time")
 
     @classmethod
-    async def create(cls, sandbox: Sandbox) -> "SandboxInstance":
+    async def create(cls, sandbox: Union[Sandbox, Dict[str, Any]]) -> "SandboxInstance":
+        if isinstance(sandbox, dict):
+            sandbox = Sandbox.from_dict(sandbox)
+
         if not sandbox.spec:
             raise Exception("Sandbox spec is required")
         if not sandbox.spec.runtime:
@@ -92,3 +103,49 @@ class SandboxInstance:
             client=client,
         )
         return response
+
+    @classmethod
+    async def create_if_not_exists(
+        cls, sandbox: Union[Sandbox, Dict[str, Any]]
+    ) -> "SandboxInstance":
+        """Create a sandbox if it doesn't exist, otherwise return existing."""
+        if isinstance(sandbox, dict):
+            sandbox = Sandbox.from_dict(sandbox)
+
+        try:
+            sandbox_instance = await cls.get(sandbox.metadata.name)
+            return sandbox_instance
+        except Exception as e:
+            # Check if it's a 404 error (sandbox not found)
+            if hasattr(e, "status_code") and e.status_code == 404:
+                return await cls.create(sandbox)
+            raise e
+
+    @classmethod
+    async def from_session(
+        cls, session: Union[SessionWithToken, Dict[str, Any]]
+    ) -> "SandboxInstance":
+        """Create a sandbox instance from a session with token."""
+        if isinstance(session, dict):
+            session = SessionWithToken.from_dict(session)
+
+        # Create a minimal sandbox configuration for session-based access
+        sandbox_name = session.name.split("-")[0] if "-" in session.name else session.name
+        sandbox = Sandbox(metadata=Metadata(name=sandbox_name))
+        config = SandboxConfiguration(
+            sandbox=sandbox,
+            force_url=session.url,
+            headers={"X-Blaxel-Preview-Token": session.token},
+            params={"bl_preview_token": session.token},
+        )
+
+        instance = cls.__new__(cls)
+        instance.sandbox = sandbox
+        instance.config = config
+        instance.fs = SandboxFileSystem(config)
+        instance.process = SandboxProcess(config)
+        instance.previews = SandboxPreviews(sandbox)
+        instance.sessions = SandboxSessions(config)
+        instance.network = SandboxNetwork(config)
+
+        return instance

blaxel/core/sandbox/session.py

@@ -0,0 +1,124 @@
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Optional, Union
+
+from ..client.api.compute import (
+    create_sandbox_preview,
+    delete_sandbox_preview,
+    get_sandbox_preview,
+    list_sandbox_preview_tokens,
+    list_sandbox_previews,
+)
+from ..client.client import client
+from ..client.models import Metadata, Preview, PreviewSpec
+from .preview import SandboxPreview
+from .types import SandboxConfiguration, SessionCreateOptions, SessionWithToken
+
+
+class SandboxSessions:
+    def __init__(self, sandbox_config: SandboxConfiguration):
+        self.sandbox_config = sandbox_config
+
+    @property
+    def sandbox_name(self) -> str:
+        return self.sandbox_config.metadata.name if self.sandbox_config.metadata else ""
+
+    async def create(
+        self, options: Optional[Union[SessionCreateOptions, Dict[str, Any]]] = None
+    ) -> SessionWithToken:
+        if options is None:
+            options = SessionCreateOptions()
+        elif isinstance(options, dict):
+            options = SessionCreateOptions.from_dict(options)
+
+        expires_at = options.expires_at or datetime.now() + timedelta(days=1)
+
+        preview_body = Preview(
+            metadata=Metadata(name=f"session-{int(datetime.now().timestamp() * 1000)}"),
+            spec=PreviewSpec(
+                port=443,
+                public=False,
+                request_headers=options.request_headers,
+                response_headers=options.response_headers,
+            ),
+        )
+
+        preview_response = await create_sandbox_preview.asyncio(
+            self.sandbox_name, client=client, body=preview_body
+        )
+
+        preview = SandboxPreview(preview_response, self.sandbox_name)
+        token_obj = await preview.tokens.create(expires_at)
+
+        return SessionWithToken(
+            name=preview_body.metadata.name,
+            url=preview.spec.url or "",
+            token=token_obj.value,
+            expires_at=token_obj.expires_at,
+        )
+
+    async def create_if_expired(
+        self,
+        options: Optional[Union[SessionCreateOptions, Dict[str, Any]]] = None,
+        delta_seconds: int = 3600,  # 1 hour
+    ) -> SessionWithToken:
+        if options is None:
+            options = SessionCreateOptions()
+        elif isinstance(options, dict):
+            options = SessionCreateOptions.from_dict(options)
+
+        all_sessions = await self.list()
+        now = datetime.now()
+        threshold = now + timedelta(seconds=delta_seconds)
+
+        if all_sessions:
+            session_data = all_sessions[0]
+            if session_data.expires_at < threshold:
+                await self.delete(session_data.name)
+                session_data = await self.create(options)
+        else:
+            session_data = await self.create(options)
+
+        return session_data
+
+    async def list(self) -> List[SessionWithToken]:
+        previews_response = await list_sandbox_previews.asyncio(self.sandbox_name, client=client)
+
+        sessions = []
+        for preview in previews_response:
+            if preview.metadata and preview.metadata.name and "session-" in preview.metadata.name:
+                token = await self.get_token(preview.metadata.name)
+                if token:
+                    sessions.append(
+                        SessionWithToken(
+                            name=preview.metadata.name,
+                            url=preview.spec.url or "",
+                            token=token.spec.token or "",
+                            expires_at=token.spec.expires_at or datetime.now(),
+                        )
+                    )
+
+        return sessions
+
+    async def get(self, name: str) -> SessionWithToken:
+        preview_response = await get_sandbox_preview.asyncio(self.sandbox_name, name, client=client)
+
+        token = await self.get_token(name)
+
+        return SessionWithToken(
+            name=name,
+            url=preview_response.spec.url or "",
+            token=token.spec.token or "" if token else "",
+            expires_at=token.spec.expires_at or datetime.now() if token else datetime.now(),
+        )
+
+    async def delete(self, name: str):
+        return await delete_sandbox_preview.asyncio(self.sandbox_name, name, client=client)
+
+    async def get_token(self, preview_name: str):
+        tokens_response = await list_sandbox_preview_tokens.asyncio(
+            self.sandbox_name, preview_name, client=client
+        )
+
+        if not tokens_response:
+            return None
+        return tokens_response[0]

blaxel/core/sandbox/types.py

@@ -0,0 +1,103 @@
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+from ..client.models import Sandbox
+
+
+class SessionCreateOptions:
+    def __init__(
+        self,
+        expires_at: Optional[datetime] = None,
+        response_headers: Optional[Dict[str, str]] = None,
+        request_headers: Optional[Dict[str, str]] = None,
+    ):
+        self.expires_at = expires_at
+        self.response_headers = response_headers or {}
+        self.request_headers = request_headers or {}
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SessionCreateOptions":
+        expires_at = None
+        if "expires_at" in data and data["expires_at"]:
+            if isinstance(data["expires_at"], str):
+                expires_at = datetime.fromisoformat(data["expires_at"].replace("Z", "+00:00"))
+            elif isinstance(data["expires_at"], datetime):
+                expires_at = data["expires_at"]
+
+        return cls(
+            expires_at=expires_at,
+            response_headers=data.get("response_headers"),
+            request_headers=data.get("request_headers"),
+        )
+
+
+class SessionWithToken:
+    def __init__(self, name: str, url: str, token: str, expires_at: datetime):
+        self.name = name
+        self.url = url
+        self.token = token
+        self.expires_at = expires_at
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SessionWithToken":
+        expires_at = data["expires_at"]
+        if isinstance(expires_at, str):
+            expires_at = datetime.fromisoformat(expires_at.replace("Z", "+00:00"))
+
+        return cls(
+            name=data["name"],
+            url=data["url"],
+            token=data["token"],
+            expires_at=expires_at,
+        )
+
+
+class SandboxConfiguration:
+    def __init__(
+        self,
+        sandbox: Sandbox,
+        force_url: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
+        params: Optional[Dict[str, str]] = None,
+    ):
+        self.sandbox = sandbox
+        self.force_url = force_url
+        self.headers = headers or {}
+        self.params = params or {}
+
+    @property
+    def metadata(self):
+        return self.sandbox.metadata
+
+    @property
+    def status(self):
+        return self.sandbox.status
+
+    @property
+    def spec(self):
+        return self.sandbox.spec
+
+
+class WatchEvent:
+    def __init__(self, op: str, path: str, name: str, content: Optional[str] = None):
+        self.op = op
+        self.path = path
+        self.name = name
+        self.content = content
+
+
+class SandboxFilesystemFile:
+    def __init__(self, path: str, content: str):
+        self.path = path
+        self.content = content
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SandboxFilesystemFile":
+        return cls(data["path"], data["content"])
+
+
+class CopyResponse:
+    def __init__(self, message: str, source: str, destination: str):
+        self.message = message
+        self.source = source
+        self.destination = destination