aisolate-client 0.1.0__py3-none-any.whl

aisolate/client/__init__.py

@@ -0,0 +1,12 @@
+from .sandbox import Sandbox
+from .sandbox_manager import SandboxManager
+from .exception import SandboxError
+from .models import Execution, KernelMessage
+
+__all__ = [
+    "Sandbox",
+    "SandboxManager",
+    "SandboxError",
+    "Execution",
+    "KernelMessage",
+]

aisolate/client/exception.py

@@ -0,0 +1,50 @@
+from typing import Any, Optional
+
+
+class SandboxClientError(Exception):
+    """Base exception for Sandbox Client errors."""
+
+    pass
+
+
+class SandboxConnectionError(SandboxClientError):
+    """Raised for network connection-related errors."""
+
+    pass
+
+
+class SandboxTimeoutError(SandboxClientError):
+    """Raised when a request to the sandbox API times out."""
+
+    pass
+
+
+class SandboxAPIError(SandboxClientError):
+    """
+    Raised for errors returned by the Sandbox API (e.g., HTTP 4xx or 5xx).
+    """
+
+    def __init__(
+        self,
+        message: str,
+        status_code: Optional[int] = None,
+        response_data: Optional[Any] = None,
+    ):
+        super().__init__(message)
+        self.status_code = status_code
+        self.response_data = response_data
+
+    def __str__(self):
+        return f"API Error (Status {self.status_code}): {super().__str__()} - Response: {self.response_data}"
+
+
+class SandboxWebSocketError(SandboxClientError):
+    """Raised for errors specific to WebSocket communication."""
+
+    pass
+
+
+class SandboxError(Exception):
+    """Raised for any errors related to the Sandbox lifecycle or execution."""
+
+    pass

aisolate/client/gateway_client.py

@@ -0,0 +1,384 @@
+import random
+import json
+import asyncio
+import logging
+import uuid
+from asyncio import Queue
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    Optional,
+    TypeVar,
+    AsyncContextManager,
+)
+from urllib.parse import urljoin
+
+import httpx
+from websockets.exceptions import ConnectionClosed
+from websockets.asyncio.client import connect as ws_connect, ClientConnection
+
+from aisolate.common.codec import decode_kernel_message
+from aisolate.common.models import KernelMessage, KernelExecutionCompleted
+from .exception import (
+    SandboxAPIError,
+    SandboxConnectionError,
+    SandboxTimeoutError,
+)
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=Dict[str, Any] | None)
+
+
+class GatewayClient:
+    """
+    Thin wrapper around HTTPX + WebSockets for the Gateway API.
+
+    This client maintains a single, long-lived WebSocket connection and supports
+    two patterns of use:
+    1. As an async context manager (recommended for most cases):
+       ```
+       async with GatewayClient(base_url) as client:
+           await client.create_sandbox("my-sbx")
+       ```
+    2. As a standalone object with manual lifecycle management:
+       ```
+       client = GatewayClient(base_url)
+       try:
+           await client.start()
+           await client.create_sandbox("my-sbx")
+       finally:
+           await client.close()
+       ```
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: Optional[str] = None,
+        *,
+        timeout: float = 60.0,
+        session: Optional[httpx.AsyncClient] = None,
+        ws_connect: Optional[
+            Callable[..., AsyncContextManager[ClientConnection]]
+        ] = None,
+        ping_interval: Optional[float] = 30.0,
+        ping_timeout: Optional[float] = 3600.0,
+    ) -> None:
+        """Initializes the GatewayClient configuration."""
+        self.base_url = base_url.rstrip("/") + "/"
+        self.api_key = api_key
+        self._timeout = timeout
+        self._ping_interval = ping_interval
+        self._ping_timeout = ping_timeout
+
+        # HTTPX session state
+        self._session = session
+        self._owns_session = session is None
+        # Webscocket connection
+        self._ws_connect = ws_connect
+
+        # WebSocket connection state
+        self._ws: Optional[ClientConnection] = None
+        self._listener_task: Optional[asyncio.Task] = None
+        self._response_queues: Dict[str, Queue[KernelMessage]] = {}
+        self._is_active: bool = False
+
+    async def start(self) -> "GatewayClient":
+        """
+        Establishes the HTTP session and WebSocket connection.
+        This method is idempotent and safe to call multiple times.
+        """
+        if self._is_active:
+            return self
+
+        # 1. Set up the HTTPX session for standard API calls
+        if self._session is None:
+            headers = {"Authorization": f"Token {self.api_key}"} if self.api_key else {}
+            self._session = httpx.AsyncClient(
+                base_url=self.base_url, timeout=self._timeout, headers=headers
+            )
+
+        # 2. Establish the persistent WebSocket connection
+        # The `websockets` library has more robust and configurable support for
+        # keepalive pings, which is crucial for maintaining long-lived connections
+        # through intermediaries like load balancers. We will use it directly
+        # instead of the `httpx` websocket support.
+        rel_path = "/gateway/ws"
+        headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
+        ws_url = self._absolute_ws_url(rel_path)
+        ws_connect_func = self._ws_connect or ws_connect
+        ws_cm = ws_connect_func(
+            ws_url,
+            open_timeout=self._timeout,
+            additional_headers=headers,
+            ping_interval=self._ping_interval,
+            ping_timeout=self._ping_timeout,
+        )
+
+        self._ws = await ws_cm.__aenter__()
+
+        # 3. Start the background task to listen for and route all incoming messages
+        self._response_queues = {}
+        self._listener_task = asyncio.create_task(self._listen_for_messages())
+
+        self._is_active = True
+        logger.info("GatewayClient started and connected.")
+        return self
+
+    async def close(
+        self, exc_type: Any = None, exc: Any = None, tb: Any = None
+    ) -> None:
+        """
+        Closes the WebSocket connection and the HTTP session.
+        This method is idempotent and safe to call multiple times.
+        """
+        if not self._is_active:
+            return
+
+        # 1. Stop the listener task
+        if self._listener_task:
+            self._listener_task.cancel()
+            try:
+                await self._listener_task
+            except asyncio.CancelledError:
+                pass
+            self._listener_task = None
+
+        # 2. Close the WebSocket connection (and wait for it to fully shut down)
+        if self._ws:
+            ws = self._ws
+            self._ws = None
+            await ws.close()
+            # if it's a websockets.WebSocketClientProtocol
+            if hasattr(ws, "wait_closed"):
+                await ws.wait_closed()
+
+        # 3. Close the HTTPX session if we own it
+        if self._owns_session and self._session:
+            await self._session.aclose()
+            self._session = None
+
+        self._is_active = False
+        logger.info("GatewayClient closed.")
+
+    async def __aenter__(self) -> "GatewayClient":
+        """Context manager entry point. Starts the client."""
+        return await self.start()
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        """Context manager exit point. Closes the client."""
+        await self.close(exc_type, exc, tb)
+
+    async def _listen_for_messages(self):
+        """
+        Listens for all messages on the WebSocket and routes them to the
+        correct response queue based on their 'request_id'.
+        """
+        if not self._ws:
+            return
+        try:
+            async for msg in self._ws:
+                try:
+                    data = self._extract_message_data(msg)
+                    if data is None:
+                        continue
+                    request_id, kernel_msg = decode_kernel_message(data)
+                    if request_id and request_id in self._response_queues:
+                        await self._response_queues[request_id].put(kernel_msg)
+                    else:
+                        logger.warning(
+                            f"Received message with unknown request_id: {request_id}"
+                        )
+                except asyncio.CancelledError:
+                    # we were explicitly cancelled — just exit
+                    return
+                except Exception:
+                    logger.exception(
+                        "Failed to decode or route incoming WebSocket message"
+                    )
+        except ConnectionClosed as e:
+            logger.info(f"WebSocket connection closed cleanly (code={e.code}).")
+        except Exception:
+            logger.exception("Listener task terminated due to an unexpected error.")
+
+        # If the connection drops, signal an error to all pending callers
+        # by placing a sentinel value (e.g., an exception) in their queues.
+        error = SandboxConnectionError("WebSocket connection lost")
+        for queue in self._response_queues.values():
+            await queue.put(error)
+
+    async def execute_code(
+        self,
+        sandbox_id: str,
+        code: str,
+        *,
+        language: str = "python3",
+        env_vars: Optional[Dict[str, str]] = None,
+        context_id: Optional[str] = None,
+    ) -> AsyncIterator[KernelMessage]:
+        """
+        Executes code by sending a message over the persistent WebSocket.
+        It then listens on a dedicated queue for corresponding messages.
+
+        Args:
+            sandbox_id: The sandbox to execute code in
+            code: The code to execute
+            language: Programming language (only used if context_id is None)
+            env_vars: Environment variables to set
+            context_id: Optional context/notebook ID. If None, uses default context for language.
+        """
+        if not self._is_active:
+            raise SandboxConnectionError(
+                "GatewayClient is not running. Use `async with` or call `await client.start()`."
+            )
+
+        request_id = str(uuid.uuid4())
+        queue: Queue[KernelMessage | Exception] = asyncio.Queue()
+        self._response_queues[request_id] = queue
+
+        try:
+            payload_data: Dict[str, Any] = {"code": code, "env_vars": env_vars}
+
+            # Use context_id if provided, otherwise use language (creates default context)
+            if context_id:
+                payload_data["context_id"] = context_id
+            else:
+                payload_data["language"] = language
+
+            payload = {
+                "request_id": request_id,
+                "type": "execute_request",
+                "sandbox_id": sandbox_id,
+                "payload": payload_data,
+            }
+            await self._ws_send(payload)
+
+            while True:
+                msg = await queue.get()
+                if isinstance(msg, Exception):
+                    raise msg
+                yield msg
+                if isinstance(msg, KernelExecutionCompleted):
+                    break
+        finally:
+            # Clean up the queue for this specific request
+            self._response_queues.pop(request_id, None)
+
+    async def _request(
+        self,
+        method: str,
+        endpoint: str,
+        *,
+        json: Any = None,
+        retries: int = 3,
+        backoff: float = 0.5,
+        timeout: float | httpx.Timeout | None = None,
+    ) -> T:
+        """
+        Make an HTTP request using the underlying AsyncClient.
+        """
+        if not self._is_active or self._session is None:
+            raise RuntimeError("Client not entered – use `async with` context manager")
+
+        if not hasattr(self._session, method.lower()):
+            raise ValueError(f"Unsupported HTTP method: {method}")
+        url = urljoin(self.base_url, endpoint)
+        http_method: Callable[..., Any] = getattr(self._session, method.lower())
+
+        attempt = 0
+        while True:
+            try:
+                # Only pass json parameter if it's provided
+                if json is not None:
+                    resp: httpx.Response = await http_method(url, json=json)
+                else:
+                    resp: httpx.Response = await http_method(url)
+                resp.raise_for_status()
+                return resp.json() if resp.content else None
+            except (httpx.TimeoutException, httpx.RequestError) as exc:
+                if attempt == retries:
+                    raise SandboxTimeoutError(f"{method} {url} failed") from exc
+                delay = (backoff * 2**attempt) + (random.uniform(0, 1) * backoff)
+                attempt += 1
+                await asyncio.sleep(delay)
+            except httpx.HTTPStatusError as exc:
+                raise SandboxAPIError(
+                    f"{method} {url} returned {exc.response.status_code}"
+                ) from exc
+
+    async def _ws_send(self, payload: dict) -> None:
+        if getattr(self._ws, "send", None):
+            await self._ws.send(json.dumps(payload))
+        else:
+            await self._ws.send_json(payload)
+
+    def _extract_message_data(self, msg: Any) -> Optional[bytes]:
+        """Extracts the relevant bytes from a websocket message, handling different client types."""
+        if isinstance(msg, dict):  # Starlette-style
+            if any(
+                kw in (msg.get("text") or msg.get("type") or "")
+                for kw in ("websocket.close",)
+            ):
+                return None  # Sentinel for close
+            return msg.get("bytes")
+
+        # websockets-style
+        return msg.encode() if isinstance(msg, str) else msg
+
+    def _absolute_ws_url(self, path: str) -> str:
+        """Convert a relative path to a ws:// or wss:// URL based on base_url."""
+        scheme = "wss" if self.base_url.startswith("https") else "ws"
+        root = self.base_url.removeprefix("https://").removeprefix("http://")
+        return f"{scheme}://{root.rstrip('/')}{path}"
+
+    async def create_sandbox(
+        self,
+        sandbox_id: str,
+    ) -> Dict[str, Any]:
+        payload = {"sandbox_id": sandbox_id}
+        return await self._request("POST", "sandboxes", json=payload)
+
+    async def delete_sandbox(self, sandbox_id: str) -> None:
+        await self._request("DELETE", f"sandboxes/{sandbox_id}")
+
+    # Context management methods
+    async def create_context(
+        self,
+        sandbox_id: str,
+        language: str = "python3",
+        *,
+        cwd: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, Any]:
+        """Create a new context (notebook/kernel) in a sandbox."""
+        payload = {"language": language}
+        if cwd:
+            payload["cwd"] = cwd
+        if env:
+            payload["env"] = env
+        return await self._request(
+            "POST",
+            f"sandboxes/{sandbox_id}/contexts",
+            json=payload,
+        )
+
+    async def list_contexts(self, sandbox_id: str) -> list[Dict[str, Any]]:
+        """List all contexts in a sandbox."""
+        return await self._request("GET", f"sandboxes/{sandbox_id}/contexts")
+
+    async def restart_context(self, sandbox_id: str, context_id: str) -> None:
+        """Restart a specific context."""
+        await self._request(
+            "POST",
+            f"sandboxes/{sandbox_id}/contexts/{context_id}/restart",
+        )
+
+    async def delete_context(self, sandbox_id: str, context_id: str) -> None:
+        """Delete a specific context."""
+        await self._request(
+            "DELETE",
+            f"sandboxes/{sandbox_id}/contexts/{context_id}",
+        )