pulse-coding-agent 0.1.0__py3-none-any.whl

pulse/sandbox/remote/client.py

@@ -0,0 +1,389 @@
+"""Remote Sandbox Client.
+
+Implements the client side of the authenticated WebSocket protocol.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import uuid
+from collections.abc import AsyncGenerator
+from contextlib import suppress
+from typing import Any
+
+import websockets
+from websockets.protocol import State
+
+from pulse.sandbox.remote.models import (
+    ExecutionResultModel,
+    SubmitExecutionRequest,
+    SubmitExecutionResponse,
+)
+from pulse.sandbox.remote.protocol import RemoteSandboxClient
+
+logger = logging.getLogger(__name__)
+
+
+class RemoteClient(RemoteSandboxClient):
+    """Client for the Remote Sandbox Server."""
+
+    def __init__(self, endpoint_url: str, auth_token: str) -> None:
+        self.endpoint_url = endpoint_url
+        self.auth_token = auth_token
+        self._ws: websockets.WebSocketClientProtocol | None = None
+        self._results: dict[str, asyncio.Future[ExecutionResultModel]] = {}
+        self._responses: dict[str, asyncio.Future[dict[str, Any]]] = {}
+        self._queues: dict[str, asyncio.Queue[tuple[str, str] | None]] = {}
+        self._listener_task: asyncio.Task[None] | None = None
+
+    async def connect(self) -> None:
+        """Establish the WebSocket connection."""
+        if self._ws and getattr(self._ws, "state", None) is not State.CLOSED:
+            return
+
+        headers = {"Authorization": f"Bearer {self.auth_token}"}
+
+        import os
+        import ssl
+
+        ssl_context = None
+
+        tls_cert = os.environ.get("PULSE_TLS_CERT")
+        tls_key = os.environ.get("PULSE_TLS_KEY")
+        tls_ca = os.environ.get("PULSE_TLS_CA")
+
+        if self.endpoint_url.startswith("wss://"):
+            if not (tls_cert and tls_key and tls_ca):
+                raise ValueError(
+                    "mTLS certificates (PULSE_TLS_CERT, PULSE_TLS_KEY, PULSE_TLS_CA) are required for wss://"
+                )
+            try:
+                ssl_context = ssl.create_default_context(
+                    ssl.Purpose.SERVER_AUTH, cafile=tls_ca
+                )
+                ssl_context.load_cert_chain(certfile=tls_cert, keyfile=tls_key)
+            except (ssl.SSLError, OSError) as e:
+                raise RuntimeError(f"Failed to load mTLS certificates: {e}")
+        elif not (
+            "127.0.0.1" in self.endpoint_url
+            or "localhost" in self.endpoint_url
+            or "[::1]" in self.endpoint_url
+        ):
+            raise RuntimeError(
+                "Insecure ws:// connections are only allowed for loopback (127.0.0.1)."
+            )
+
+        self._ws = await websockets.connect(
+            self.endpoint_url,
+            additional_headers=headers,
+            ping_interval=20,
+            ping_timeout=20,
+            open_timeout=5,
+            close_timeout=5,
+            ssl=ssl_context,
+        )
+        self._listener_task = asyncio.create_task(self._listen())
+
+    async def disconnect(self) -> None:
+        """Close the WebSocket connection and reap its listener task."""
+        if self._ws:
+            await self._ws.close()
+            self._ws = None
+        if self._listener_task:
+            listener = self._listener_task
+            self._listener_task = None
+            if not listener.done():
+                listener.cancel()
+            with suppress(asyncio.CancelledError):
+                await listener
+
+    async def _listen(self) -> None:
+        """Background task to listen for messages from the server."""
+        if not self._ws:
+            return
+
+        try:
+            async for message in self._ws:
+                try:
+                    data = json.loads(message)
+                    msg_type = data.get("type")
+                    payload = data.get("payload", {})
+
+                    if msg_type == "result":
+                        execution_id = payload.get("execution_id")
+                        fut = self._results.get(execution_id) if execution_id else None
+                        if fut and not fut.done():
+                            fut.set_result(ExecutionResultModel.from_dict(payload))
+                        elif execution_id:
+                            logger.warning(
+                                "Received result for unknown execution %s.",
+                                execution_id,
+                            )
+
+                        queue = self._queues.get(execution_id or "")
+                        if queue:
+                            await queue.put(None)
+
+                    elif msg_type == "response":
+                        request_id = payload.get("request_id")
+                        if request_id and request_id in self._responses:
+                            fut = self._responses[request_id]
+                            if not fut.done():
+                                fut.set_result(payload)
+                            continue
+                        eid = payload.get("execution_id")
+                        if eid and eid in self._responses:
+                            fut = self._responses[eid]
+                            if not fut.done():
+                                fut.set_result(payload)
+                        else:
+                            logger.warning(
+                                "Received response without a known request ID: %s",
+                                payload,
+                            )
+
+                    elif msg_type == "stream":
+                        stream_data = payload.get("data", "")
+                        stream_type = payload.get("stream", "stdout")
+                        execution_id = payload.get("execution_id")
+                        queue = self._queues.get(execution_id or "")
+                        if queue:
+                            await queue.put((stream_type, stream_data))
+
+                    elif msg_type == "error":
+                        request_id = (
+                            payload.get("request_id")
+                            if isinstance(payload, dict)
+                            else None
+                        )
+                        error = (
+                            payload.get("message", "Remote server error")
+                            if isinstance(payload, dict)
+                            else str(payload)
+                        )
+                        fut = self._responses.get(request_id) if request_id else None
+                        if fut and not fut.done():
+                            fut.set_exception(RuntimeError(error))
+                        else:
+                            logger.error("Remote server error: %s", error)
+
+                except (KeyError, ValueError, RuntimeError) as e:
+                    # check if the message is a direct string payload from attach history
+                    if (
+                        isinstance(message, str)
+                        and message.startswith("{")
+                        and '"stream"' in message
+                    ):
+                        try:
+                            line_data = json.loads(message)
+                            if line_data.get("type") == "stream":
+                                stream_payload = line_data.get("payload", {})
+                                stream_data = stream_payload.get("data", "")
+                                stream_type = stream_payload.get("stream", "stdout")
+                                for queue in self._queues.values():
+                                    await queue.put((stream_type, stream_data))
+                        except (TypeError, ValueError):
+                            logger.error(f"Error processing server message: {e}")
+                    else:
+                        logger.error(f"Error processing server message: {e}")
+        except websockets.exceptions.ConnectionClosed:
+            logger.info("Connection to remote server closed.")
+            for fut in self._results.values():
+                if not fut.done():
+                    fut.set_exception(RuntimeError("Connection closed"))
+            for fut in self._responses.values():
+                if not fut.done():
+                    fut.set_exception(RuntimeError("Connection closed"))
+            for queue in self._queues.values():
+                await queue.put(None)
+
+    async def submit(self, request: SubmitExecutionRequest) -> SubmitExecutionResponse:
+        """Submit an execution request to the remote worker."""
+        await self.connect()
+
+        self._results[request.execution_id] = asyncio.get_running_loop().create_future()
+        self._responses[request.execution_id] = (
+            asyncio.get_running_loop().create_future()
+        )
+        self._queues[request.execution_id] = asyncio.Queue()
+
+        payload = request.to_dict()
+        payload["request_id"] = request.execution_id
+        await self._ws.send(
+            json.dumps(
+                {
+                    "action": "submit",
+                    "payload": payload,
+                }
+            )
+        )
+
+        resp = await self._responses[request.execution_id]
+        return SubmitExecutionResponse(
+            execution_id=request.execution_id, status=resp.get("status", "STARTING")
+        )
+
+    async def cancel(self, execution_id: str) -> None:
+        """Cancel an ongoing execution on the remote worker."""
+        if not self._ws or getattr(self._ws, "state", None) is State.CLOSED:
+            return
+
+        await self._ws.send(
+            json.dumps({"action": "cancel", "payload": {"execution_id": execution_id}})
+        )
+
+    async def get_result(self, execution_id: str) -> ExecutionResultModel:
+        """Wait for and retrieve the final execution result."""
+        fut = self._results.get(execution_id)
+        if not fut:
+            raise ValueError(f"No active execution found for {execution_id}")
+
+        try:
+            return await fut
+        finally:
+            self._results.pop(execution_id, None)
+
+    async def stream_output(
+        self, execution_id: str
+    ) -> AsyncGenerator[tuple[str, str], None]:
+        """Stream stdout and stderr from the remote execution."""
+        queue = self._queues.get(execution_id)
+        if not queue:
+            return
+
+        try:
+            while True:
+                item = await queue.get()
+                if item is None:
+                    break
+                yield item
+        finally:
+            self._queues.pop(execution_id, None)
+
+    async def status(self, execution_id: str) -> str:
+        """Query authoritative remote state."""
+        await self.connect()
+        fut = asyncio.get_running_loop().create_future()
+        temp_key = f"status-{uuid.uuid4().hex}"
+        self._responses[temp_key] = fut
+        try:
+            await self._ws.send(
+                json.dumps(
+                    {
+                        "action": "status",
+                        "payload": {
+                            "execution_id": execution_id,
+                            "request_id": temp_key,
+                        },
+                    }
+                )
+            )
+            resp = await fut
+            return resp.get("status", "NOT_FOUND")
+        finally:
+            self._responses.pop(temp_key, None)
+
+    async def attach(self, execution_id: str) -> ExecutionResultModel:
+        """Attach to a live execution and stream results, returning the final result."""
+        await self.connect()
+
+        # Setup tracking
+        self._results[execution_id] = asyncio.get_running_loop().create_future()
+        request_id = f"attach-{uuid.uuid4().hex}"
+        self._responses[request_id] = asyncio.get_running_loop().create_future()
+        self._queues[execution_id] = asyncio.Queue()
+
+        await self._ws.send(
+            json.dumps(
+                {
+                    "action": "attach",
+                    "payload": {"execution_id": execution_id, "request_id": request_id},
+                }
+            )
+        )
+
+        fut = self._results.get(execution_id)
+        if not fut:
+            raise ValueError(f"No active execution found for {execution_id}")
+
+        try:
+            response = await self._responses[request_id]
+            if response.get("status") not in {"ATTACHED", "COMPLETED", "FAILED"}:
+                raise RuntimeError(
+                    f"Remote execution {execution_id} cannot be attached: "
+                    f"{response.get('status', 'UNKNOWN')}"
+                )
+            return await fut
+        finally:
+            self._results.pop(execution_id, None)
+            self._responses.pop(request_id, None)
+            self._queues.pop(execution_id, None)
+
+    async def reconcile(self) -> None:
+        """Reconcile orphaned or stale executions with the remote worker."""
+        await self.connect()
+        fut = asyncio.get_running_loop().create_future()
+        temp_key = f"reconcile-{uuid.uuid4().hex}"
+        self._responses[temp_key] = fut
+        try:
+            await self._ws.send(
+                json.dumps({"action": "reconcile", "payload": {"request_id": temp_key}})
+            )
+            await fut
+        finally:
+            self._responses.pop(temp_key, None)
+
+    async def upload_artifact(self, execution_id: str, archive_data: bytes) -> None:
+        """Upload a workspace snapshot to the remote worker before execution."""
+        import base64
+
+        await self.connect()
+        fut = asyncio.get_running_loop().create_future()
+        # Use a temporary key for the response
+        temp_key = f"upload_{execution_id}"
+        self._responses[temp_key] = fut
+        try:
+            await self._ws.send(
+                json.dumps(
+                    {
+                        "action": "upload_artifact",
+                        "payload": {
+                            "execution_id": execution_id,
+                            "request_id": temp_key,
+                            "data": base64.b64encode(archive_data).decode("ascii"),
+                        },
+                    }
+                )
+            )
+            await fut
+        finally:
+            self._responses.pop(temp_key, None)
+
+    async def download_artifact(self, execution_id: str) -> bytes:
+        """Download the modified workspace overlay from the remote worker after execution."""
+        import base64
+
+        await self.connect()
+        fut = asyncio.get_running_loop().create_future()
+        temp_key = f"download_{execution_id}"
+        self._responses[temp_key] = fut
+        try:
+            await self._ws.send(
+                json.dumps(
+                    {
+                        "action": "download_artifact",
+                        "payload": {
+                            "execution_id": execution_id,
+                            "request_id": temp_key,
+                        },
+                    }
+                )
+            )
+            resp = await fut
+            if resp.get("status") == "DOWNLOADED" and "data" in resp:
+                return base64.b64decode(resp["data"])
+            return b""
+        finally:
+            self._responses.pop(temp_key, None)

pulse/sandbox/remote/models.py

@@ -0,0 +1,167 @@
+"""Data models for Remote Sandbox Protocol."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import PurePosixPath
+from typing import Any
+
+REMOTE_PROTOCOL_VERSION = "1.0"
+_SAFE_IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")
+_SAFE_ENV_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,127}$")
+
+
+def validate_execution_id(value: object) -> str:
+    if not isinstance(value, str) or not _SAFE_IDENTIFIER.fullmatch(value):
+        raise ValueError("execution_id must be a safe 1-128 character identifier")
+    return value
+
+
+@dataclass
+class SubmitExecutionRequest:
+    protocol_version: str
+    execution_id: str
+    idempotency_key: str
+    command: str | list[str]
+    correlation_id: str | None = None
+    working_directory: str | None = None
+    env: dict[str, str] | None = None
+    resource_policy_dict: dict[str, Any] | None = None
+    network_policy_dict: dict[str, Any] | None = None
+    secret_policy_dict: dict[str, Any] | None = None
+
+    def __post_init__(self) -> None:
+        if self.protocol_version != REMOTE_PROTOCOL_VERSION:
+            raise ValueError("unsupported remote protocol version")
+        validate_execution_id(self.execution_id)
+        if not _SAFE_IDENTIFIER.fullmatch(self.idempotency_key):
+            raise ValueError("idempotency_key must be a safe 1-128 character identifier")
+        if isinstance(self.command, str):
+            if not self.command or len(self.command) > 65_536 or "\x00" in self.command:
+                raise ValueError("command must be 1-65536 characters without NUL bytes")
+        elif isinstance(self.command, list):
+            if not self.command or len(self.command) > 256 or any(
+                not isinstance(item, str)
+                or not item
+                or len(item) > 65_536
+                or "\x00" in item
+                for item in self.command
+            ):
+                raise ValueError("command arguments are invalid or exceed protocol limits")
+        else:
+            raise TypeError("command must be a string or an array of strings")
+        if self.working_directory:
+            path = PurePosixPath(self.working_directory.replace("\\", "/"))
+            if path.is_absolute() or ".." in path.parts:
+                raise ValueError("working_directory must remain inside the execution workspace")
+        if self.env is not None:
+            if not isinstance(self.env, dict) or len(self.env) > 128:
+                raise ValueError("env must be an object with at most 128 entries")
+            if any(
+                not isinstance(key, str)
+                or not _SAFE_ENV_NAME.fullmatch(key)
+                or not isinstance(value, str)
+                or len(value) > 65_536
+                or "\x00" in value
+                for key, value in self.env.items()
+            ):
+                raise ValueError("env contains an invalid name or value")
+        for policy in (
+            self.resource_policy_dict,
+            self.network_policy_dict,
+            self.secret_policy_dict,
+        ):
+            if policy is not None and not isinstance(policy, dict):
+                raise ValueError("execution policies must be objects")
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "protocol_version": self.protocol_version,
+            "execution_id": self.execution_id,
+            "idempotency_key": self.idempotency_key,
+            "command": self.command,
+            "correlation_id": self.correlation_id,
+            "working_directory": self.working_directory,
+            "env": self.env,
+            "resource_policy_dict": self.resource_policy_dict,
+            "network_policy_dict": self.network_policy_dict,
+            "secret_policy_dict": self.secret_policy_dict,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> SubmitExecutionRequest:
+        return cls(
+            protocol_version=data["protocol_version"],
+            execution_id=data["execution_id"],
+            idempotency_key=data["idempotency_key"],
+            command=data["command"],
+            correlation_id=data.get("correlation_id"),
+            working_directory=data.get("working_directory"),
+            env=data.get("env"),
+            resource_policy_dict=data.get("resource_policy_dict"),
+            network_policy_dict=data.get("network_policy_dict"),
+            secret_policy_dict=data.get("secret_policy_dict"),
+        )
+
+
+@dataclass
+class SubmitExecutionResponse:
+    execution_id: str
+    status: str
+    error: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "execution_id": self.execution_id,
+            "status": self.status,
+            "error": self.error,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> SubmitExecutionResponse:
+        return cls(
+            execution_id=data["execution_id"],
+            status=data["status"],
+            error=data.get("error"),
+        )
+
+
+@dataclass
+class ExecutionResultModel:
+    execution_id: str
+    command: str
+    exit_code: int
+    stdout: str
+    stderr: str
+    duration_ms: float
+    timed_out: bool = False
+    truncated: bool = False
+    termination_reason: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "execution_id": self.execution_id,
+            "command": self.command,
+            "exit_code": self.exit_code,
+            "stdout": self.stdout,
+            "stderr": self.stderr,
+            "duration_ms": self.duration_ms,
+            "timed_out": self.timed_out,
+            "truncated": self.truncated,
+            "termination_reason": self.termination_reason,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ExecutionResultModel:
+        return cls(
+            execution_id=data["execution_id"],
+            command=data["command"],
+            exit_code=data["exit_code"],
+            stdout=data["stdout"],
+            stderr=data["stderr"],
+            duration_ms=data["duration_ms"],
+            timed_out=data.get("timed_out", False),
+            truncated=data.get("truncated", False),
+            termination_reason=data.get("termination_reason"),
+        )

pulse/sandbox/remote/protocol.py

@@ -0,0 +1,65 @@
+"""Remote Sandbox Protocol definition."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from typing import Protocol
+
+from pulse.sandbox.remote.models import (
+    ExecutionResultModel,
+    SubmitExecutionRequest,
+    SubmitExecutionResponse,
+)
+
+
+class RemoteSandboxClient(Protocol):
+    """Protocol for communicating with a Remote Sandbox Worker.
+
+    The implementation must support:
+    - Authenticated transport
+    - Artifact snapshot staging and retrieval
+    - Execution submission
+    - Execution cancellation
+    - Execution reconciliation
+    """
+
+    async def submit(self, request: SubmitExecutionRequest) -> SubmitExecutionResponse:
+        """Submit an execution request to the remote worker."""
+        ...
+
+    async def cancel(self, execution_id: str) -> None:
+        """Cancel an ongoing execution on the remote worker."""
+        ...
+
+    async def get_result(self, execution_id: str) -> ExecutionResultModel:
+        """Wait for and retrieve the final execution result."""
+        ...
+
+    async def stream_output(
+        self, execution_id: str
+    ) -> AsyncGenerator[tuple[str, str], None]:
+        """Stream stdout and stderr from the remote execution.
+
+        Yields tuples of (stdout_chunk, stderr_chunk).
+        """
+        ...
+
+    async def status(self, execution_id: str) -> str:
+        """Query authoritative execution state (RUNNING, COMPLETED, FAILED, NOT_FOUND)."""
+        ...
+
+    async def attach(self, execution_id: str) -> ExecutionResultModel:
+        """Attach to a live execution and stream results, returning the final result."""
+        ...
+
+    async def reconcile(self) -> None:
+        """Reconcile orphaned or stale executions with the remote worker."""
+        ...
+
+    async def upload_artifact(self, execution_id: str, archive_data: bytes) -> None:
+        """Upload a workspace snapshot to the remote worker before execution."""
+        ...
+
+    async def download_artifact(self, execution_id: str) -> bytes:
+        """Download the modified workspace overlay from the remote worker after execution."""
+        ...