aisolate-client 0.2.0__tar.gz → 0.3.0__tar.gz

{aisolate_client-0.2.0 → aisolate_client-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aisolate-client
-Version: 0.2.0
+Version: 0.3.0
 Summary: Client library for sandbox environments
 License: Apache-2.0
 Requires-Python: >=3.10
@@ -16,7 +16,7 @@ Client library for sandbox environments
 ## Local Installation
 
 ```bash
-pip install path/to/aisolate_client-0.2.0-py3-none-any.whl
+pip install path/to/aisolate_client-0.3.0-py3-none-any.whl
 ```
 
 ## Usage

{aisolate_client-0.2.0 → aisolate_client-0.3.0}/README.md

@@ -5,7 +5,7 @@ Client library for sandbox environments
 ## Local Installation
 
 ```bash
-pip install path/to/aisolate_client-0.2.0-py3-none-any.whl
+pip install path/to/aisolate_client-0.3.0-py3-none-any.whl
 ```
 
 ## Usage

{aisolate_client-0.2.0 → aisolate_client-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "aisolate-client"
-version = "0.2.0"
+version = "0.3.0"
 description = "Client library for sandbox environments"
 readme = "README.md"
 requires-python = ">=3.10"

{aisolate_client-0.2.0 → aisolate_client-0.3.0}/src/aisolate/client/gateway_client.py

@@ -13,13 +13,22 @@ from typing import (
     TypeVar,
     AsyncContextManager,
 )
-from urllib.parse import urljoin
+from urllib.parse import quote, 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.codec import (
+    MAX_MULTIPLEX_FRAME_BYTES,
+    decode_frame,
+    decode_multiplex_frame,
+)
+
+#: Per-sandbox bound on the client's live change feed. A chatty (or hostile)
+#: sandbox can only fill its own queue; drop-oldest on overflow keeps the shared
+#: client's memory bounded, and the /workspace/changes resync recovers the gap.
+CHANGES_QUEUE_MAXSIZE = 1024
 from aisolate.common.models import KernelMessage, KernelExecutionCompleted
 from .exception import (
     SandboxAPIError,
@@ -60,6 +69,7 @@ class GatewayClient:
         api_key: Optional[str] = None,
         *,
         timeout: float = 60.0,
+        connect_timeout: Optional[float] = None,
         session: Optional[httpx.AsyncClient] = None,
         ws_connect: Optional[
             Callable[..., AsyncContextManager[ClientConnection]]
@@ -67,10 +77,19 @@ class GatewayClient:
         ping_interval: Optional[float] = 30.0,
         ping_timeout: Optional[float] = 3600.0,
     ) -> None:
-        """Initializes the GatewayClient configuration."""
+        """Initializes the GatewayClient configuration.
+
+        ``timeout`` is the blanket read/write/pool timeout of the httpx
+        session. ``connect_timeout`` (when set) bounds connection
+        establishment separately — including the WebSocket open — so a
+        long read timeout (e.g. a create request parked in the gateway's
+        admission queue) does not also slow down failure detection
+        against an unreachable or hung host.
+        """
         self.base_url = base_url.rstrip("/") + "/"
         self.api_key = api_key
         self._timeout = timeout
+        self._connect_timeout = connect_timeout
         self._ping_interval = ping_interval
         self._ping_timeout = ping_timeout
 
@@ -83,7 +102,13 @@ class GatewayClient:
         # WebSocket connection state
         self._ws: Optional[ClientConnection] = None
         self._listener_task: Optional[asyncio.Task] = None
+        # Per-request exec response queues, keyed by (globally unique) request_id.
         self._response_queues: Dict[str, Queue[KernelMessage]] = {}
+        # request_id → sandbox_id, so a per-sandbox failure fails only that
+        # sandbox's in-flight requests, never every co-tenant's.
+        self._request_sandbox: Dict[str, str] = {}
+        # Per-sandbox live change feeds (channel="changes"), bounded.
+        self._changes_queues: Dict[str, Queue[Any]] = {}
         self._is_active: bool = False
 
     async def start(self) -> "GatewayClient":
@@ -97,8 +122,13 @@ class GatewayClient:
         # 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 {}
+            timeout = (
+                httpx.Timeout(self._timeout, connect=self._connect_timeout)
+                if self._connect_timeout is not None
+                else self._timeout
+            )
             self._session = httpx.AsyncClient(
-                base_url=self.base_url, timeout=self._timeout, headers=headers
+                base_url=self.base_url, timeout=timeout, headers=headers
             )
 
         # 2. Establish the persistent WebSocket connection
@@ -112,17 +142,26 @@ class GatewayClient:
         ws_connect_func = self._ws_connect or ws_connect
         ws_cm = ws_connect_func(
             ws_url,
-            open_timeout=self._timeout,
+            open_timeout=(
+                self._connect_timeout
+                if self._connect_timeout is not None
+                else self._timeout
+            ),
             additional_headers=headers,
             ping_interval=self._ping_interval,
             ping_timeout=self._ping_timeout,
-            max_size=10 * 1024 * 1024,  # 10 MB – prevent 1009 on large kernel output
+            # Defence in depth: the gateway already drops backend frames that
+            # would exceed this, scoped to one sandbox, so this cap is never the
+            # thing that closes the shared connection.
+            max_size=MAX_MULTIPLEX_FRAME_BYTES,
         )
 
         self._ws = await ws_cm.__aenter__()
 
         # 3. Start the background task to listen for and route all incoming messages
         self._response_queues = {}
+        self._request_sandbox = {}
+        self._changes_queues = {}
         self._listener_task = asyncio.create_task(self._listen_for_messages())
 
         self._is_active = True
@@ -173,50 +212,113 @@ class GatewayClient:
         """Context manager exit point. Closes the client."""
         await self.close(exc_type, exc, tb)
 
-    async def _listen_for_messages(self):
+    def _changes_queue_for(self, sandbox_id: str) -> "Queue[Any]":
+        """The bounded live-change feed for one sandbox, created on demand."""
+        q = self._changes_queues.get(sandbox_id)
+        if q is None:
+            q = asyncio.Queue(maxsize=CHANGES_QUEUE_MAXSIZE)
+            self._changes_queues[sandbox_id] = q
+        return q
+
+    def _push_change(self, sandbox_id: str, change: Any) -> None:
+        """Enqueue a change for a sandbox, dropping the oldest if the feed is full.
+
+        Bounding per sandbox means a flooding sandbox can only fill its own
+        queue; the dropped record is recovered by the /workspace/changes resync.
         """
-        Listens for all messages on the WebSocket and routes them to the
-        correct response queue based on their 'request_id'.
+        q = self._changes_queue_for(sandbox_id)
+        if q.full():
+            try:
+                q.get_nowait()
+            except asyncio.QueueEmpty:
+                pass
+        q.put_nowait(change)
+
+    def _fail_sandbox_requests(self, sandbox_id: str, message: str) -> None:
+        """Fail only the in-flight exec requests belonging to one sandbox."""
+        error = SandboxConnectionError(message)
+        for request_id, sid in list(self._request_sandbox.items()):
+            if sid == sandbox_id:
+                queue = self._response_queues.get(request_id)
+                if queue is not None:
+                    queue.put_nowait(error)
+
+    def _fail_all_requests(self, message: str) -> None:
+        """Fail every in-flight request — the whole connection is gone."""
+        error = SandboxConnectionError(message)
+        for queue in self._response_queues.values():
+            queue.put_nowait(error)
+
+    def _handle_sandbox_control(self, sandbox_id: str, control: Dict[str, Any]) -> None:
+        """Act on a gateway control signal scoped to one sandbox."""
+        if control.get("control") == "frame_dropped":
+            logger.warning(
+                "Gateway dropped a frame for sandbox %s (%s, %s bytes); "
+                "failing only its requests",
+                sandbox_id,
+                control.get("reason"),
+                control.get("size"),
+            )
+            self._fail_sandbox_requests(
+                sandbox_id, f"sandbox frame dropped: {control.get('reason')}"
+            )
+
+    async def _listen_for_messages(self):
+        """Route every inbound frame, isolating failures to one sandbox.
+
+        Each frame is wrapped by the gateway with its source ``sandbox_id``. A
+        frame that fails to decode fails only that sandbox's requests and the
+        listener keeps running for every co-tenant — one sandbox can never take
+        the shared connection (or another sandbox's work) down with it. Only a
+        real connection loss fails everyone.
         """
         if not self._ws:
             return
         try:
             async for msg in self._ws:
+                data = self._extract_message_data(msg)
+                if data is None:
+                    continue
                 try:
-                    data = self._extract_message_data(msg)
-                    if data is None:
+                    sandbox_id, payload = decode_multiplex_frame(data)
+                except Exception:
+                    # Can't attribute it to a sandbox — drop this frame only and
+                    # keep listening; never tear the loop down over one frame.
+                    logger.exception("Failed to peel multiplex envelope; dropping frame")
+                    continue
+                try:
+                    if isinstance(payload, dict):
+                        self._handle_sandbox_control(sandbox_id, payload)
                         continue
-                    request_id, kernel_msg = decode_kernel_message(data)
-                    if request_id and request_id in self._response_queues:
+                    channel, request_id, kernel_msg = decode_frame(payload)
+                    if channel == "changes":
+                        # Unsolicited push, correlated by cursor, not request_id.
+                        self._push_change(sandbox_id, kernel_msg)
+                    elif 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}"
+                            "Unknown request_id %s for sandbox %s",
+                            request_id,
+                            sandbox_id,
                         )
                 except asyncio.CancelledError:
-                    # we were explicitly cancelled — just exit
                     return
                 except Exception:
                     logger.exception(
-                        "Failed to decode or route incoming WebSocket message"
+                        "Failed to decode/route frame; failing only sandbox %s",
+                        sandbox_id,
                     )
-                    # Propagate decode errors to all pending callers so they
-                    # don't hang forever waiting for KernelExecutionCompleted.
-                    error = SandboxConnectionError(
-                        "Failed to decode sandbox message"
+                    self._fail_sandbox_requests(
+                        sandbox_id, "Failed to decode sandbox message"
                     )
-                    for queue in self._response_queues.values():
-                        await queue.put(error)
         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)
+        # The connection itself dropped: every sandbox's requests are dead.
+        self._fail_all_requests("WebSocket connection lost")
 
     async def execute_code(
         self,
@@ -246,6 +348,7 @@ class GatewayClient:
         request_id = str(uuid.uuid4())
         queue: Queue[KernelMessage | Exception] = asyncio.Queue()
         self._response_queues[request_id] = queue
+        self._request_sandbox[request_id] = sandbox_id
 
         try:
             payload_data: Dict[str, Any] = {"code": code, "env_vars": env_vars}
@@ -274,6 +377,7 @@ class GatewayClient:
         finally:
             # Clean up the queue for this specific request
             self._response_queues.pop(request_id, None)
+            self._request_sandbox.pop(request_id, None)
 
     async def _request(
         self,
@@ -297,6 +401,7 @@ class GatewayClient:
         http_method: Callable[..., Any] = getattr(self._session, method.lower())
 
         attempt = 0
+        retried_503 = False
         while True:
             try:
                 # Only pass json parameter if it's provided
@@ -313,6 +418,17 @@ class GatewayClient:
                 attempt += 1
                 await asyncio.sleep(delay)
             except httpx.HTTPStatusError as exc:
+                # 503 = the gateway's RAM-budget admission queue is full
+                # or timed out; it sets Retry-After. One honored retry —
+                # callers that want to fail fast still get the error on
+                # sustained pressure.
+                if exc.response.status_code == 503 and not retried_503:
+                    retried_503 = True
+                    delay = min(
+                        float(exc.response.headers.get("retry-after", 5)), 60.0
+                    )
+                    await asyncio.sleep(delay)
+                    continue
                 try:
                     response_data = exc.response.json()
                 except Exception:
@@ -360,7 +476,12 @@ class GatewayClient:
             payload["ramMb"] = ram_mb
         if vcpu is not None:
             payload["vcpu"] = vcpu
-        return await self._request("POST", "sandboxes", json=payload)
+        # No transport retries here: callers (the agent's provisioning
+        # barrier) already retry creation with a fresh connection, and an
+        # inner re-POST of the same sandbox_id after a read timeout risks
+        # double admission of a request the gateway may still be holding
+        # in its queue. The 503 Retry-After branch is unaffected.
+        return await self._request("POST", "sandboxes", json=payload, retries=0)
 
     async def delete_sandbox(self, sandbox_id: str) -> None:
         await self._request("DELETE", f"sandboxes/{sandbox_id}")
@@ -403,3 +524,54 @@ class GatewayClient:
             "DELETE",
             f"sandboxes/{sandbox_id}/contexts/{context_id}",
         )
+
+    # Workspace (files) methods
+    async def list_tree(self, sandbox_id: str, root: str = "") -> Dict[str, Any]:
+        """List the sandbox workspace tree ({root, entries}).
+
+        ``root`` restricts the walk to one workspace subtree (per-agent subdir);
+        entry paths stay workspace-root-relative either way.
+        """
+        suffix = f"?root={quote(root, safe='')}" if root else ""
+        return await self._request(
+            "GET", f"sandboxes/{sandbox_id}/workspace/tree{suffix}"
+        )
+
+    async def read_file(self, sandbox_id: str, path: str) -> Dict[str, Any]:
+        """Read one workspace file (utf-8 text or base64 binary)."""
+        return await self._request(
+            "GET",
+            f"sandboxes/{sandbox_id}/workspace/file?path={quote(path, safe='')}",
+        )
+
+    async def write_file(
+        self,
+        sandbox_id: str,
+        path: str,
+        content: str,
+        encoding: str = "utf-8",
+    ) -> Dict[str, Any]:
+        """Create or overwrite one workspace file."""
+        return await self._request(
+            "PUT",
+            f"sandboxes/{sandbox_id}/workspace/file?path={quote(path, safe='')}",
+            json={"content": content, "encoding": encoding},
+        )
+
+    async def workspace_changes(self, sandbox_id: str, since: int = 0) -> Dict[str, Any]:
+        """Return the workspace change delta since ``since`` ({cursor, changes, snapshot})."""
+        return await self._request(
+            "GET",
+            f"sandboxes/{sandbox_id}/workspace/changes?since={int(since)}",
+        )
+
+    async def aiter_workspace_changes(self, sandbox_id: str) -> AsyncIterator[Any]:
+        """Yield each pushed ``WorkspaceChange`` for one sandbox as it arrives.
+
+        Attributed by the gateway's per-frame ``sandbox_id`` tag, so one client
+        multiplexing several sandboxes gets a separate, non-interleaved feed for
+        each — a chatty or hostile sandbox cannot pollute another's stream.
+        """
+        q = self._changes_queue_for(sandbox_id)
+        while True:
+            yield await q.get()

{aisolate_client-0.2.0 → aisolate_client-0.3.0}/src/aisolate/client/models.py

@@ -12,6 +12,7 @@ from aisolate.common.models import (
     ErrorOutput,
     ExecuteResult,
     KernelMessage,
+    ShellExitStatus,
     StreamOutput,
 )
 
@@ -124,27 +125,6 @@ class Execution:
         """Returns True if the execution produced a final answer."""
         return self.final_answer is not None
 
-    @property
-    def review_request(self) -> Optional[Dict[str, Any]]:
-        """Returns the review request data if present in DisplayData.
-
-        Checks for the custom MIME type 'application/vnd.aisolate.review-request+json'
-        which is sent by the request_report_review() action.
-        """
-        for event in self.events:
-            if isinstance(event, DisplayData):
-                review_data = event.data.get(
-                    "application/vnd.aisolate.review-request+json"
-                )
-                if review_data:
-                    return review_data
-        return None
-
-    @property
-    def has_review_request(self) -> bool:
-        """Returns True if the execution requested a review."""
-        return self.review_request is not None
-
     def _find_main_result_event(
         self,
     ) -> Optional[Union["ExecuteResult", "DisplayData"]]:
@@ -320,6 +300,27 @@ class Execution:
             if isinstance(e, StreamOutput) and e.name == "stdout"
         )
 
+    @property
+    def stderr(self) -> str:
+        """Concatenates all text from standard error streams."""
+        return "".join(
+            e.text
+            for e in self.events
+            if isinstance(e, StreamOutput) and e.name == "stderr"
+        )
+
+    @property
+    def exit_code(self) -> Optional[int]:
+        """The shell command's process exit code, if this was a shell run.
+
+        Carried by the ShellExecutor's ``shell_exit`` message; ``None`` for
+        kernel executions (which signal errors via ``ErrorOutput`` instead).
+        """
+        return next(
+            (e.exit_code for e in self.events if isinstance(e, ShellExitStatus)),
+            None,
+        )
+
     def __repr__(self) -> str:
         if self.has_error:
             status = f"error='{self.error.error_name}'"