From 42f2ab96891c0ac848b9d8aad55f27ec147b1e90 Mon Sep 17 00:00:00 2001 From: Dhiraj Bokde Date: Thu, 16 Jul 2026 13:19:53 -0700 Subject: [PATCH] feat(python): add async client (grpc.aio) Add asyncio-native sandbox, workspace, and inference-route clients that mirror the current synchronous Python SDK using grpc.aio. Preserve workspace identity across sandbox CRUD, sessions, waits, and high-level context management. Reuse the existing TLS and OIDC refresh plumbing, including async interceptors for every RPC shape. Document the async surface and add unit coverage for execution, authentication, workspace routing, inference routes, lifecycle operations, and context management. Signed-off-by: Dhiraj Bokde --- docs/sandboxes/manage-sandboxes.mdx | 23 + python/openshell/__init__.py | 10 + python/openshell/async_sandbox_test.py | 934 ++++++++++++++++++++++ python/openshell/sandbox.py | 1004 ++++++++++++++++++++++-- 4 files changed, 1918 insertions(+), 53 deletions(-) create mode 100644 python/openshell/async_sandbox_test.py diff --git a/docs/sandboxes/manage-sandboxes.mdx b/docs/sandboxes/manage-sandboxes.mdx index bc408c4ecd..788bb060e5 100644 --- a/docs/sandboxes/manage-sandboxes.mdx +++ b/docs/sandboxes/manage-sandboxes.mdx @@ -242,6 +242,29 @@ with SandboxClient.from_active_cluster() as client: assert sandbox.id in {s.id for s in matches} ``` +## Use the Async Python Client + +Use `AsyncSandboxClient` in asyncio applications to avoid blocking the event +loop. It mirrors the synchronous sandbox lifecycle and execution methods with +`async`/`await`, uses `grpc.aio`, and requires the same explicit workspace +selection: + +```python +from openshell import AsyncSandboxClient + +async with AsyncSandboxClient.from_active_cluster() as client: + sandbox = await client.create(workspace="default", name="async-agent") + result = await client.exec(sandbox.id, ["python", "-c", "print('ready')"]) + print(result.stdout) + await client.delete(sandbox.name, workspace=sandbox.workspace) +``` + +`AsyncSandbox` provides the corresponding high-level async context manager. +`AsyncWorkspaceClient` and `AsyncInferenceRouteClient` expose async workspace +lifecycle and workspace-scoped inference route operations. Authentication, +TLS, OIDC refresh, command timeouts, and streaming execution follow the same +semantics as the synchronous clients. + ## Expose Long Running Services Service forwarding makes a long-running process inside a sandbox reachable through a gateway-managed URL. Use it for development servers, notebooks, dashboards, or other services that keep listening after the sandbox starts. Run the service on loopback inside the sandbox, expose its port, then open the URL printed by OpenShell. diff --git a/python/openshell/__init__.py b/python/openshell/__init__.py index 0aa343ed60..6ad8eb950f 100644 --- a/python/openshell/__init__.py +++ b/python/openshell/__init__.py @@ -6,6 +6,11 @@ from __future__ import annotations from .sandbox import ( + AsyncInferenceRouteClient, + AsyncSandbox, + AsyncSandboxClient, + AsyncSandboxSession, + AsyncWorkspaceClient, ExecChunk, ExecResult, InferenceRouteClient, @@ -29,6 +34,11 @@ __version__ = "0.0.0" __all__ = [ + "AsyncInferenceRouteClient", + "AsyncSandbox", + "AsyncSandboxClient", + "AsyncSandboxSession", + "AsyncWorkspaceClient", "ExecChunk", "ExecResult", "InferenceRouteClient", diff --git a/python/openshell/async_sandbox_test.py b/python/openshell/async_sandbox_test.py new file mode 100644 index 0000000000..1496ba35f7 --- /dev/null +++ b/python/openshell/async_sandbox_test.py @@ -0,0 +1,934 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, cast + +import grpc +import grpc.aio +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +from openshell._proto import datamodel_pb2, openshell_pb2 +from openshell.sandbox import ( + _PYTHON_CLOUDPICKLE_BOOTSTRAP, + _SANDBOX_PYTHON_BIN, + AsyncInferenceRouteClient, + AsyncSandbox, + AsyncSandboxClient, + AsyncSandboxSession, + AsyncWorkspaceClient, + ExecResult, + SandboxError, + SandboxRef, + SandboxStatusRef, + _AsyncBearerUnaryStreamInterceptor, + _AsyncBearerUnaryUnaryInterceptor, +) + + +async def _aiter(items: list[Any]) -> Any: + for item in items: + yield item + + +def _async_client_with_fake_stub(stub: object) -> AsyncSandboxClient: + client = cast("AsyncSandboxClient", object.__new__(AsyncSandboxClient)) + client._timeout = 30.0 + client._stub = cast("Any", stub) + return client + + +def _make_sandbox_proto( + id_: str, + name: str, + labels: dict[str, str] | None = None, + phase: openshell_pb2.SandboxPhase = openshell_pb2.SANDBOX_PHASE_READY, + version: int = 0, + workspace: str = "default", +) -> openshell_pb2.Sandbox: + sandbox = openshell_pb2.Sandbox() + sandbox.metadata.id = id_ + sandbox.metadata.name = name + sandbox.metadata.workspace = workspace + for key, value in (labels or {}).items(): + sandbox.metadata.labels[key] = value + sandbox.status.phase = phase + sandbox.status.current_policy_version = version + return sandbox + + +class _FakeAsyncExecStub: + def __init__(self, events: list[openshell_pb2.ExecSandboxEvent] | None = None): + self.request: openshell_pb2.ExecSandboxRequest | None = None + self._events = events or [ + openshell_pb2.ExecSandboxEvent( + exit=openshell_pb2.ExecSandboxExit(exit_code=0) + ) + ] + + def ExecSandbox( + self, + request: openshell_pb2.ExecSandboxRequest, + timeout: float | None = None, + ) -> Any: + self.request = request + _ = timeout + return _aiter(list(self._events)) + + +class _FakeAsyncInferenceStub: + def __init__(self) -> None: + self.request: Any = None + + async def SetInferenceRoute( + self, request: Any, timeout: float | None = None + ) -> Any: + self.request = request + _ = timeout + return SimpleNamespace( + provider_name=request.provider_name, + model_id=request.model_id, + version=1, + ) + + async def GetInferenceRoute( + self, request: Any, timeout: float | None = None + ) -> Any: + self.request = request + _ = timeout + return SimpleNamespace( + provider_name="openai-dev", + model_id="gpt-4.1", + version=2, + ) + + async def DeleteInferenceRoute( + self, request: Any, timeout: float | None = None + ) -> Any: + self.request = request + _ = timeout + return SimpleNamespace(deleted=True) + + +class _FakeAsyncSandboxStub: + def __init__(self, listed: list[openshell_pb2.Sandbox] | None = None) -> None: + self.create_request: openshell_pb2.CreateSandboxRequest | None = None + self.list_request: openshell_pb2.ListSandboxesRequest | None = None + self.delete_request: openshell_pb2.DeleteSandboxRequest | None = None + self.health_calls = 0 + self._listed = listed or [] + + async def CreateSandbox( + self, + request: openshell_pb2.CreateSandboxRequest, + timeout: float | None = None, + ) -> Any: + self.create_request = request + _ = timeout + return SimpleNamespace( + sandbox=_make_sandbox_proto( + "sandbox-1", + request.name or "generated", + dict(request.labels), + workspace=request.workspace, + ) + ) + + async def ListSandboxes( + self, + request: openshell_pb2.ListSandboxesRequest, + timeout: float | None = None, + ) -> Any: + self.list_request = request + _ = timeout + return SimpleNamespace(sandboxes=list(self._listed)) + + async def DeleteSandbox( + self, + request: openshell_pb2.DeleteSandboxRequest, + timeout: float | None = None, + ) -> Any: + self.delete_request = request + _ = timeout + return SimpleNamespace(deleted=True) + + async def Health(self, request: Any, timeout: float | None = None) -> Any: + self.health_calls += 1 + _ = (request, timeout) + return openshell_pb2.HealthResponse() + + +class _FakeAsyncWorkspaceStub: + def __init__(self) -> None: + self.request: Any = None + + @staticmethod + def _workspace(name: str, labels: dict[str, str] | None = None) -> Any: + workspace = datamodel_pb2.Workspace() + workspace.metadata.name = name + workspace.metadata.labels.update(labels or {}) + workspace.status.phase = datamodel_pb2.WORKSPACE_PHASE_ACTIVE + return workspace + + async def CreateWorkspace(self, request: Any, timeout: float | None = None) -> Any: + self.request = request + _ = timeout + return SimpleNamespace( + workspace=self._workspace(request.name, dict(request.labels)) + ) + + async def ListWorkspaces(self, request: Any, timeout: float | None = None) -> Any: + self.request = request + _ = timeout + return SimpleNamespace(workspaces=[self._workspace("staging")]) + + async def DeleteWorkspace(self, request: Any, timeout: float | None = None) -> Any: + self.request = request + _ = timeout + return SimpleNamespace(deleted=True) + + +# --------------------------------------------------------------------------- +# exec / exec_python +# --------------------------------------------------------------------------- + + +async def test_async_exec_sends_stdin_payload() -> None: + stub = _FakeAsyncExecStub() + client = _async_client_with_fake_stub(stub) + + result = await client.exec( + "sandbox-1", ["python", "-c", "print('ok')"], stdin=b"payload" + ) + + assert result.exit_code == 0 + assert stub.request is not None + assert stub.request.stdin == b"payload" + + +async def test_async_exec_python_serializes_callable_payload() -> None: + stub = _FakeAsyncExecStub() + client = _async_client_with_fake_stub(stub) + + def add(a: int, b: int) -> int: + return a + b + + result = await client.exec_python("sandbox-1", add, args=(2, 3)) + + assert result.exit_code == 0 + assert stub.request is not None + assert stub.request.command == [ + _SANDBOX_PYTHON_BIN, + "-c", + _PYTHON_CLOUDPICKLE_BOOTSTRAP, + ] + assert stub.request.environment["OPENSHELL_PYFUNC_B64"] + assert stub.request.stdin == b"" + + +async def test_async_exec_stream_yields_chunks_then_result() -> None: + stub = _FakeAsyncExecStub( + events=[ + openshell_pb2.ExecSandboxEvent( + stdout=openshell_pb2.ExecSandboxStdout(data=b"out") + ), + openshell_pb2.ExecSandboxEvent( + stderr=openshell_pb2.ExecSandboxStderr(data=b"err") + ), + openshell_pb2.ExecSandboxEvent( + exit=openshell_pb2.ExecSandboxExit(exit_code=7) + ), + ] + ) + client = _async_client_with_fake_stub(stub) + + items = [item async for item in client.exec_stream("sandbox-1", ["ls"])] + + result = items[-1] + assert isinstance(result, ExecResult) + assert result.exit_code == 7 + assert result.stdout == "out" + assert result.stderr == "err" + + +async def test_async_exec_stream_rejects_empty_command() -> None: + client = _async_client_with_fake_stub(_FakeAsyncExecStub()) + + with pytest.raises(SandboxError, match="command must not be empty"): + [item async for item in client.exec_stream("sandbox-1", [])] + + +# --------------------------------------------------------------------------- +# Bearer auth interceptor (async twin) +# --------------------------------------------------------------------------- + + +def test_async_bearer_interceptor_attaches_authorization_header() -> None: + interceptor = _AsyncBearerUnaryUnaryInterceptor(lambda: "secret-token") + details = grpc.aio.ClientCallDetails( + method="/Test/Method", + timeout=None, + metadata=grpc.aio.Metadata(("x-existing", "yes")), + credentials=None, + wait_for_ready=None, + ) + + new_details = interceptor._attach(details) + assert new_details.metadata is not None + md = list(new_details.metadata) + + assert ("x-existing", "yes") in md + assert ("authorization", "Bearer secret-token") in md + + +def test_async_bearer_interceptor_handles_empty_metadata() -> None: + interceptor = _AsyncBearerUnaryStreamInterceptor(lambda: "t") + details = grpc.aio.ClientCallDetails( + method="/Test/Method", + timeout=None, + metadata=None, + credentials=None, + wait_for_ready=None, + ) + + new_details = interceptor._attach(details) + + assert new_details.metadata is not None + assert list(new_details.metadata) == [("authorization", "Bearer t")] + + +async def test_async_bearer_interceptor_awaits_continuation() -> None: + interceptor = _AsyncBearerUnaryUnaryInterceptor(lambda: "tok") + captured: dict[str, Any] = {} + + async def continuation(details: Any, request: Any) -> str: + captured["details"] = details + captured["request"] = request + return "result" + + details = grpc.aio.ClientCallDetails( + method="/Test/Method", + timeout=None, + metadata=None, + credentials=None, + wait_for_ready=None, + ) + result = await interceptor.intercept_unary_unary(continuation, details, "payload") + + assert result == "result" + assert ("authorization", "Bearer tok") in list(captured["details"].metadata) + assert captured["request"] == "payload" + + +def test_async_bearer_interceptor_calls_provider_per_request() -> None: + tokens = iter(["t1", "t2", "t3"]) + interceptor = _AsyncBearerUnaryUnaryInterceptor(lambda: next(tokens)) + seen: list[str] = [] + + for _ in range(3): + details = grpc.aio.ClientCallDetails( + method="/m", + timeout=None, + metadata=None, + credentials=None, + wait_for_ready=None, + ) + attached = interceptor._attach(details) + assert attached.metadata is not None + for key, value in attached.metadata: + if key == "authorization": + assert isinstance(value, str) + seen.append(value) + + assert seen == ["Bearer t1", "Bearer t2", "Bearer t3"] + + +# --------------------------------------------------------------------------- +# from_active_cluster: gateway resolution shared with the sync client +# --------------------------------------------------------------------------- + + +def _setup_gateway_dir( + tmp_path: Path, + monkeypatch: Any, + *, + name: str = "g", + endpoint: str = "http://127.0.0.1:8080", + auth_mode: str | None = None, + mtls_files: dict[str, str] | None = None, + oidc_bundle: dict | None = None, +) -> Path: + gateway_dir = tmp_path / "openshell" / "gateways" / name + gateway_dir.mkdir(parents=True) + (tmp_path / "openshell" / "active_gateway").write_text(name) + meta: dict[str, Any] = {"gateway_endpoint": endpoint} + if auth_mode is not None: + meta["auth_mode"] = auth_mode + (gateway_dir / "metadata.json").write_text(json.dumps(meta)) + if mtls_files: + mtls_dir = gateway_dir / "mtls" + mtls_dir.mkdir() + for fname, body in mtls_files.items(): + (mtls_dir / fname).write_text(body) + if oidc_bundle is not None: + (gateway_dir / "oidc_token.json").write_text(json.dumps(oidc_bundle)) + monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path)) + monkeypatch.delenv("OPENSHELL_GATEWAY", raising=False) + return gateway_dir + + +def _bearer_interceptor_count(channel: Any) -> int: + """Count bearer interceptors across all four grpc.aio call categories. + + grpc.aio sorts interceptors into per-call-type buckets, so an OIDC + client should register exactly one interceptor in each of the four.""" + total = 0 + for attr in ( + "_unary_unary_interceptors", + "_unary_stream_interceptors", + "_stream_unary_interceptors", + "_stream_stream_interceptors", + ): + total += len(getattr(channel, attr, [])) + return total + + +async def test_async_from_active_cluster_reads_gateway_metadata_layout( + tmp_path: Path, + monkeypatch: Any, +) -> None: + gateway_name = "test-gateway" + gateway_dir = tmp_path / "openshell" / "gateways" / gateway_name + mtls_dir = gateway_dir / "mtls" + mtls_dir.mkdir(parents=True) + (tmp_path / "openshell" / "active_gateway").write_text(gateway_name) + (gateway_dir / "metadata.json").write_text( + json.dumps({"gateway_endpoint": "https://127.0.0.1:8443"}) + ) + (mtls_dir / "ca.crt").write_text("ca") + (mtls_dir / "tls.crt").write_text("cert") + (mtls_dir / "tls.key").write_text("key") + + monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path)) + monkeypatch.delenv("OPENSHELL_GATEWAY", raising=False) + + client = AsyncSandboxClient.from_active_cluster() + try: + assert client._cluster_name == gateway_name + assert client._endpoint == "127.0.0.1:8443" + finally: + await client.close() + + +async def test_async_from_active_cluster_prefers_openshell_gateway_env( + tmp_path: Path, + monkeypatch: Any, +) -> None: + gateway_name = "env-gateway" + gateway_dir = tmp_path / "openshell" / "gateways" / gateway_name + mtls_dir = gateway_dir / "mtls" + mtls_dir.mkdir(parents=True) + (gateway_dir / "metadata.json").write_text( + json.dumps({"gateway_endpoint": "https://127.0.0.1:8443"}) + ) + (mtls_dir / "ca.crt").write_text("ca") + (mtls_dir / "tls.crt").write_text("cert") + (mtls_dir / "tls.key").write_text("key") + + monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path)) + monkeypatch.setenv("OPENSHELL_GATEWAY", gateway_name) + + client = AsyncSandboxClient.from_active_cluster() + try: + assert client._cluster_name == gateway_name + finally: + await client.close() + + +async def test_async_from_active_cluster_loads_bearer_when_auth_mode_is_oidc( + tmp_path: Path, + monkeypatch: Any, +) -> None: + _setup_gateway_dir( + tmp_path, + monkeypatch, + auth_mode="oidc", + oidc_bundle={"access_token": "from-disk"}, + ) + client = AsyncSandboxClient.from_active_cluster() + try: + # One bearer interceptor per call category so streaming exec is + # authenticated too. + assert _bearer_interceptor_count(client._channel) == 4 + finally: + await client.close() + + +async def test_async_from_active_cluster_ignores_stale_token_when_not_oidc( + tmp_path: Path, + monkeypatch: Any, +) -> None: + _setup_gateway_dir( + tmp_path, + monkeypatch, + oidc_bundle={"access_token": "stale-from-disk"}, + ) + client = AsyncSandboxClient.from_active_cluster() + try: + assert _bearer_interceptor_count(client._channel) == 0 + finally: + await client.close() + + +async def test_async_from_active_cluster_https_oidc_without_mtls_uses_tls( + tmp_path: Path, + monkeypatch: Any, +) -> None: + _setup_gateway_dir( + tmp_path, + monkeypatch, + endpoint="https://gateway.example:443", + auth_mode="oidc", + oidc_bundle={"access_token": "t"}, + ) + client = AsyncSandboxClient.from_active_cluster() + try: + assert client._endpoint == "gateway.example:443" + assert _bearer_interceptor_count(client._channel) == 4 + finally: + await client.close() + + +# --------------------------------------------------------------------------- +# CRUD surface: create / list / delete / health +# --------------------------------------------------------------------------- + + +async def test_async_create_forwards_name_and_labels() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + + ref = await client.create( + workspace="default", + name="job-1", + labels={"aiq": "deep-research"}, + ) + + assert stub.create_request is not None + assert stub.create_request.name == "job-1" + assert stub.create_request.workspace == "default" + assert dict(stub.create_request.labels) == {"aiq": "deep-research"} + assert dict(ref.labels) == {"aiq": "deep-research"} + + +async def test_async_create_without_args_sends_empty_metadata() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + + await client.create(workspace="default") + + assert stub.create_request is not None + assert stub.create_request.name == "" + assert dict(stub.create_request.labels) == {} + + +async def test_async_create_copies_caller_labels() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + + caller_labels = {"aiq": "deep-research"} + await client.create(workspace="default", labels=caller_labels) + caller_labels["aiq"] = "mutated" + + assert stub.create_request is not None + assert dict(stub.create_request.labels) == {"aiq": "deep-research"} + + +async def test_async_create_session_forwards_name_and_labels() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + + session = await client.create_session( + workspace="default", name="job-2", labels={"team": "aiq"} + ) + + assert isinstance(session, AsyncSandboxSession) + assert stub.create_request is not None + assert stub.create_request.name == "job-2" + assert dict(stub.create_request.labels) == {"team": "aiq"} + assert session.sandbox.name == "job-2" + + +async def test_async_session_delete_preserves_workspace() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + session = await client.create_session(workspace="staging", name="job-2") + + assert await session.delete() is True + assert stub.delete_request is not None + assert stub.delete_request.workspace == "staging" + + +async def test_async_list_forwards_label_selector() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + + await client.list(workspace="default", label_selector="aiq=deep-research") + + assert stub.list_request is not None + assert stub.list_request.workspace == "default" + assert stub.list_request.label_selector == "aiq=deep-research" + + +async def test_async_list_without_selector_sends_empty_string() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + + await client.list(workspace="default") + + assert stub.list_request is not None + assert stub.list_request.label_selector == "" + + +async def test_async_list_for_all_workspaces_sets_flag() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + + await client.list_for_all_workspaces() + + assert stub.list_request is not None + assert stub.list_request.all_workspaces is True + assert stub.list_request.workspace == "" + + +async def test_async_list_ids_forwards_label_selector() -> None: + stub = _FakeAsyncSandboxStub(listed=[_make_sandbox_proto("sandbox-1", "job-1")]) + client = _async_client_with_fake_stub(stub) + + ids = await client.list_ids(workspace="default", label_selector="aiq=deep-research") + + assert stub.list_request is not None + assert stub.list_request.label_selector == "aiq=deep-research" + assert ids == ["sandbox-1"] + + +async def test_async_delete_returns_bool() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + + assert await client.delete("job-1", workspace="default") is True + assert stub.delete_request is not None + assert stub.delete_request.name == "job-1" + assert stub.delete_request.workspace == "default" + + +async def test_async_health_calls_stub() -> None: + stub = _FakeAsyncSandboxStub() + client = _async_client_with_fake_stub(stub) + + await client.health() + + assert stub.health_calls == 1 + + +# --------------------------------------------------------------------------- +# wait_ready / wait_deleted +# --------------------------------------------------------------------------- + + +class _ReadyStub: + def __init__(self, phase: openshell_pb2.SandboxPhase) -> None: + self._phase = phase + + async def GetSandbox(self, request: Any, timeout: float | None = None) -> Any: + _ = timeout + return SimpleNamespace( + sandbox=_make_sandbox_proto("sandbox-1", request.name, phase=self._phase) + ) + + +async def test_async_wait_ready_returns_when_ready() -> None: + client = _async_client_with_fake_stub(_ReadyStub(openshell_pb2.SANDBOX_PHASE_READY)) + + ref = await client.wait_ready("job-1", workspace="default", timeout_seconds=5) + + assert ref.status.phase == openshell_pb2.SANDBOX_PHASE_READY + + +async def test_async_wait_ready_raises_on_error_phase() -> None: + client = _async_client_with_fake_stub(_ReadyStub(openshell_pb2.SANDBOX_PHASE_ERROR)) + + with pytest.raises(SandboxError, match="error phase"): + await client.wait_ready("job-1", workspace="default", timeout_seconds=5) + + +class _NotFoundStub: + async def GetSandbox(self, request: Any, timeout: float | None = None) -> Any: + _ = (request, timeout) + raise grpc.aio.AioRpcError( + grpc.StatusCode.NOT_FOUND, + grpc.aio.Metadata(), + grpc.aio.Metadata(), + details="not found", + ) + + +async def test_async_wait_deleted_returns_on_not_found() -> None: + client = _async_client_with_fake_stub(_NotFoundStub()) + + # Should return promptly (no timeout) when the sandbox is already gone. + await client.wait_deleted("job-1", workspace="default", timeout_seconds=5) + + +# --------------------------------------------------------------------------- +# Inference route client (async twin) +# --------------------------------------------------------------------------- + + +async def test_async_inference_set_route_forwards_workspace_and_no_verify() -> None: + stub = _FakeAsyncInferenceStub() + client = cast( + "AsyncInferenceRouteClient", object.__new__(AsyncInferenceRouteClient) + ) + client._timeout = 30.0 + client._stub = cast("Any", stub) + + await client.set_route( + workspace="staging", + provider_name="openai-dev", + model_id="gpt-4.1", + no_verify=True, + ) + + assert stub.request is not None + assert stub.request.workspace == "staging" + assert stub.request.no_verify is True + + +async def test_async_inference_get_and_delete_route_forward_workspace() -> None: + stub = _FakeAsyncInferenceStub() + client = cast( + "AsyncInferenceRouteClient", object.__new__(AsyncInferenceRouteClient) + ) + client._timeout = 30.0 + client._stub = cast("Any", stub) + + route = await client.get_route(workspace="staging") + assert route.version == 2 + assert stub.request.workspace == "staging" + + assert await client.delete_route(workspace="staging", route_name="primary") + assert stub.request.workspace == "staging" + assert stub.request.route_name == "primary" + + +async def test_async_workspace_client_mirrors_lifecycle_operations() -> None: + stub = _FakeAsyncWorkspaceStub() + client = cast("AsyncWorkspaceClient", object.__new__(AsyncWorkspaceClient)) + client._timeout = 30.0 + client._stub = cast("Any", stub) + + created = await client.create("staging", labels={"team": "agents"}) + assert created.name == "staging" + assert created.labels == {"team": "agents"} + + listed = await client.list(label_selector="team=agents") + assert [workspace.name for workspace in listed] == ["staging"] + assert stub.request.label_selector == "team=agents" + + assert await client.delete("staging") is True + assert stub.request.name == "staging" + + +# --------------------------------------------------------------------------- +# Lifecycle: context manager, close, bearer_close +# --------------------------------------------------------------------------- + + +async def test_async_client_context_manager_returns_self() -> None: + async with AsyncSandboxClient("localhost:8080") as client: + assert isinstance(client, AsyncSandboxClient) + + +async def test_async_client_close_invokes_bearer_close() -> None: + closed = [0] + + def bearer_close() -> None: + closed[0] += 1 + + client = AsyncSandboxClient( + "localhost:8080", + bearer_token="tok", + _bearer_close=bearer_close, + ) + await client.close() + assert closed[0] == 1 + # close() is idempotent — re-invoking does not double-call. + await client.close() + assert closed[0] == 1 + + +# --------------------------------------------------------------------------- +# High-level AsyncSandbox context manager +# --------------------------------------------------------------------------- + + +class _RecordingAsyncClient: + def __init__(self) -> None: + self.create_kwargs: dict[str, Any] | None = None + self.closed = False + + async def create_session( + self, + *, + workspace: str, + spec: Any = None, + name: str | None = None, + labels: Any = None, + ) -> Any: + self.create_kwargs = { + "workspace": workspace, + "spec": spec, + "name": name, + "labels": labels, + } + return SimpleNamespace( + sandbox=SimpleNamespace(name=name or "generated"), + _workspace=workspace, + ) + + async def wait_ready( + self, + name: str, + *, + workspace: str, + timeout_seconds: float = 300.0, + ) -> SandboxRef: + _ = timeout_seconds + return SandboxRef( + id="sandbox-1", + name=name, + workspace=workspace, + status=SandboxStatusRef(phase=2, current_policy_version=0), + ) + + async def close(self) -> None: + self.closed = True + + +async def test_async_sandbox_wrapper_forwards_auth_kwargs(monkeypatch: Any) -> None: + captured: dict[str, Any] = {} + + class _Sentinel(Exception): + pass + + def fake_from_active_cluster(**kwargs: Any) -> Any: + captured.update(kwargs) + raise _Sentinel + + monkeypatch.setattr( + AsyncSandboxClient, + "from_active_cluster", + staticmethod(fake_from_active_cluster), + ) + + sandbox = AsyncSandbox( + workspace="default", + cluster="my-gw", + timeout=42.0, + auto_refresh=False, + write_back=False, + insecure=True, + ) + + with pytest.raises(_Sentinel): + await sandbox.__aenter__() + + assert captured["cluster"] == "my-gw" + assert captured["timeout"] == 42.0 + assert captured["auto_refresh"] is False + assert captured["write_back"] is False + assert captured["insecure"] is True + + +async def test_async_sandbox_wrapper_defaults_match(monkeypatch: Any) -> None: + captured: dict[str, Any] = {} + + class _Sentinel(Exception): + pass + + def fake_from_active_cluster(**kwargs: Any) -> Any: + captured.update(kwargs) + raise _Sentinel + + monkeypatch.setattr( + AsyncSandboxClient, + "from_active_cluster", + staticmethod(fake_from_active_cluster), + ) + + with pytest.raises(_Sentinel): + await AsyncSandbox(workspace="default").__aenter__() + + assert captured["auto_refresh"] is True + assert captured["write_back"] is True + assert captured["insecure"] is False + + +async def test_async_high_level_creation_forwards_name_and_labels( + monkeypatch: Any, +) -> None: + recording = _RecordingAsyncClient() + monkeypatch.setattr( + AsyncSandboxClient, + "from_active_cluster", + staticmethod(lambda **_kwargs: recording), + ) + + sandbox = AsyncSandbox( + workspace="staging", + name="job-1", + labels={"aiq": "deep-research"}, + delete_on_exit=False, + ) + await sandbox.__aenter__() + try: + assert recording.create_kwargs == { + "workspace": "staging", + "spec": None, + "name": "job-1", + "labels": {"aiq": "deep-research"}, + } + finally: + await sandbox.__aexit__(None, None, None) + assert recording.closed is True + + +async def test_async_high_level_attach_rejects_name() -> None: + sandbox = AsyncSandbox( + workspace="default", sandbox="existing-sandbox", name="job-1" + ) + + with pytest.raises(SandboxError): + await sandbox.__aenter__() + + +async def test_async_high_level_attach_rejects_labels() -> None: + ref = SandboxRef( + id="sandbox-1", + name="existing", + workspace="default", + status=SandboxStatusRef(phase=2, current_policy_version=0), + ) + sandbox = AsyncSandbox( + workspace="default", sandbox=ref, labels={"aiq": "deep-research"} + ) + + with pytest.raises(SandboxError): + await sandbox.__aenter__() diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index b62a9c8855..b53d9354f8 100644 --- a/python/openshell/sandbox.py +++ b/python/openshell/sandbox.py @@ -3,6 +3,7 @@ from __future__ import annotations +import asyncio import base64 import contextlib import json @@ -18,6 +19,7 @@ from urllib.parse import urlparse import grpc +import grpc.aio import httpx from ._proto import ( @@ -40,7 +42,13 @@ class _ClientCallDetails(_ClientCallDetailsBase, grpc.ClientCallDetails): if TYPE_CHECKING: import builtins - from collections.abc import Callable, Iterator, Mapping, Sequence + from collections.abc import ( + AsyncIterator, + Callable, + Iterator, + Mapping, + Sequence, + ) @dataclass(frozen=True) @@ -348,60 +356,19 @@ def from_active_cluster( CLI's `--insecure` flag for issuers behind self-signed certs. Off by default. """ - cluster_name = cluster or _resolve_active_cluster() - gateway_dir = _xdg_config_home() / "openshell" / "gateways" / cluster_name - metadata_path = gateway_dir / "metadata.json" - try: - metadata = json.loads(metadata_path.read_text(encoding="utf-8")) - except FileNotFoundError: - raise SandboxError(f"gateway '{cluster_name}' not found") from None - if "gateway_endpoint" not in metadata: - raise SandboxError(f"gateway '{cluster_name}' metadata missing endpoint") - parsed = urlparse(metadata["gateway_endpoint"]) - host = parsed.hostname or "127.0.0.1" - port = parsed.port or (443 if parsed.scheme == "https" else 80) - endpoint = f"{host}:{port}" - - # TLS transport. Mirror crates/openshell-tui/src/lib.rs - # `build_oidc_channel` — for an https gateway, always build a - # secure channel and pick the strongest available trust profile. - tls: TlsConfig | None = None - if parsed.scheme == "https": - mtls_dir = gateway_dir / "mtls" - ca = mtls_dir / "ca.crt" if (mtls_dir / "ca.crt").exists() else None - cert = mtls_dir / "tls.crt" if (mtls_dir / "tls.crt").exists() else None - key = mtls_dir / "tls.key" if (mtls_dir / "tls.key").exists() else None - if ca is not None and cert is not None and key is not None: - # Full mTLS. - tls = TlsConfig(ca_path=ca, cert_path=cert, key_path=key) - elif ca is not None: - # CA-only trust (no client identity). - tls = TlsConfig(ca_path=ca) - else: - # System roots (e.g. OIDC gateway behind a public CA). - tls = TlsConfig() - - # OIDC bearer. Mirror the Rust CLI/TUI: the gateway metadata's - # `auth_mode` is authoritative — a stale oidc_token.json next to - # a non-OIDC gateway should NOT cause us to attach a bearer. - bearer_token: Callable[[], str] | None = None - bearer_close: Callable[[], None] | None = None - if metadata.get("auth_mode") == "oidc": - bearer_token, bearer_close = _make_cluster_bearer_provider( - gateway_dir, - cluster_name, - auto_refresh=auto_refresh, - write_back=write_back, - insecure=insecure, - ) - + conn = _resolve_gateway_connection( + cluster=cluster, + auto_refresh=auto_refresh, + write_back=write_back, + insecure=insecure, + ) return cls( - endpoint, - tls=tls, - bearer_token=bearer_token, + conn.endpoint, + tls=conn.tls, + bearer_token=conn.bearer_token, timeout=timeout, - cluster_name=cluster_name, - _bearer_close=bearer_close, + cluster_name=conn.cluster_name, + _bearer_close=conn.bearer_close, ) def close(self) -> None: @@ -1608,3 +1575,934 @@ def _resolve_active_cluster() -> str: if value == "": raise SandboxError("no active gateway configured") return value + + +@dataclass(frozen=True) +class _GatewayConnection: + """Resolved transport + auth material for an active gateway. + + Shared by the sync and async clients' `from_active_cluster` so both + derive endpoint, TLS, and OIDC bearer wiring from the identical + on-disk gateway state (see `_resolve_gateway_connection`). + """ + + endpoint: str + tls: TlsConfig | None + bearer_token: Callable[[], str] | None + bearer_close: Callable[[], None] | None + cluster_name: str + + +def _resolve_gateway_connection( + *, + cluster: str | None, + auto_refresh: bool, + write_back: bool, + insecure: bool, +) -> _GatewayConnection: + """Read an active gateway's on-disk state into connection material. + + Factored out of `SandboxClient.from_active_cluster` so the async + `AsyncSandboxClient.from_active_cluster` resolves endpoint, TLS + trust profile, and OIDC bearer provider from the same gateway + directory layout. The OIDC kwargs (`auto_refresh`, `write_back`, + `insecure`) carry the semantics documented on + `SandboxClient.from_active_cluster`. + """ + cluster_name = cluster or _resolve_active_cluster() + gateway_dir = _xdg_config_home() / "openshell" / "gateways" / cluster_name + metadata_path = gateway_dir / "metadata.json" + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + except FileNotFoundError: + raise SandboxError(f"gateway '{cluster_name}' not found") from None + if "gateway_endpoint" not in metadata: + raise SandboxError(f"gateway '{cluster_name}' metadata missing endpoint") + parsed = urlparse(metadata["gateway_endpoint"]) + host = parsed.hostname or "127.0.0.1" + port = parsed.port or (443 if parsed.scheme == "https" else 80) + endpoint = f"{host}:{port}" + + # TLS transport. Mirror crates/openshell-tui/src/lib.rs + # `build_oidc_channel` — for an https gateway, always build a + # secure channel and pick the strongest available trust profile. + tls: TlsConfig | None = None + if parsed.scheme == "https": + mtls_dir = gateway_dir / "mtls" + ca = mtls_dir / "ca.crt" if (mtls_dir / "ca.crt").exists() else None + cert = mtls_dir / "tls.crt" if (mtls_dir / "tls.crt").exists() else None + key = mtls_dir / "tls.key" if (mtls_dir / "tls.key").exists() else None + if ca is not None and cert is not None and key is not None: + # Full mTLS. + tls = TlsConfig(ca_path=ca, cert_path=cert, key_path=key) + elif ca is not None: + # CA-only trust (no client identity). + tls = TlsConfig(ca_path=ca) + else: + # System roots (e.g. OIDC gateway behind a public CA). + tls = TlsConfig() + + # OIDC bearer. Mirror the Rust CLI/TUI: the gateway metadata's + # `auth_mode` is authoritative — a stale oidc_token.json next to + # a non-OIDC gateway should NOT cause us to attach a bearer. + bearer_token: Callable[[], str] | None = None + bearer_close: Callable[[], None] | None = None + if metadata.get("auth_mode") == "oidc": + bearer_token, bearer_close = _make_cluster_bearer_provider( + gateway_dir, + cluster_name, + auto_refresh=auto_refresh, + write_back=write_back, + insecure=insecure, + ) + + return _GatewayConnection( + endpoint=endpoint, + tls=tls, + bearer_token=bearer_token, + bearer_close=bearer_close, + cluster_name=cluster_name, + ) + + +# --------------------------------------------------------------------------- +# Async client (asyncio over grpc.aio) +# +# Mirrors the synchronous surface above using `grpc.aio`. The sync client +# stays the canonical implementation; this async variant reuses the same +# proto stubs, dataclasses, TLS/bearer plumbing, and on-disk gateway state +# so the two stay in lockstep. Method names and shapes match their sync +# counterparts, differing only by `async`/`await` and `async with`. +# --------------------------------------------------------------------------- + + +class _AsyncBearerAuthMixin: + """Shared `authorization: Bearer ` attachment for `grpc.aio`. + + Async counterpart of `_BearerAuthInterceptor._attach`. The token + provider is the same synchronous per-call callable used by the sync + client (see `_normalize_bearer` and `_make_cluster_bearer_provider`), + invoked once per RPC so runtime token rotation is picked up + identically. A refresh performed by an `_OidcRefresher`-backed + provider is synchronous, matching the sync client — it runs inline + rather than on the event loop's executor. + + Unlike sync gRPC (where a single `grpc.intercept_channel` interceptor + covers every call type), `grpc.aio` sorts interceptors into per-call + categories by their interface and stops at the first match. So one + interceptor object implementing all four interfaces would only ever + fire for unary-unary calls. To attach the bearer to *every* RPC — + including the server-streaming `ExecSandbox` — this mixin is combined + with exactly one `grpc.aio` interceptor interface per subclass, and + `AsyncSandboxClient` registers one instance of each. + """ + + def __init__(self, token_provider: Callable[[], str]) -> None: + self._token_provider = token_provider + + def _attach( + self, details: grpc.aio.ClientCallDetails + ) -> grpc.aio.ClientCallDetails: + metadata = grpc.aio.Metadata() + if details.metadata is not None: + for key, value in details.metadata: + metadata.add(key, value) + metadata.add("authorization", f"Bearer {self._token_provider()}") + return grpc.aio.ClientCallDetails( + method=details.method, + timeout=details.timeout, + metadata=metadata, + credentials=details.credentials, + wait_for_ready=details.wait_for_ready, + ) + + +class _AsyncBearerUnaryUnaryInterceptor( + _AsyncBearerAuthMixin, grpc.aio.UnaryUnaryClientInterceptor +): + async def intercept_unary_unary(self, continuation, client_call_details, request): + return await continuation(self._attach(client_call_details), request) + + +class _AsyncBearerUnaryStreamInterceptor( + _AsyncBearerAuthMixin, grpc.aio.UnaryStreamClientInterceptor +): + async def intercept_unary_stream(self, continuation, client_call_details, request): + return await continuation(self._attach(client_call_details), request) + + +class _AsyncBearerStreamUnaryInterceptor( + _AsyncBearerAuthMixin, grpc.aio.StreamUnaryClientInterceptor +): + async def intercept_stream_unary( + self, continuation, client_call_details, request_iterator + ): + return await continuation(self._attach(client_call_details), request_iterator) + + +class _AsyncBearerStreamStreamInterceptor( + _AsyncBearerAuthMixin, grpc.aio.StreamStreamClientInterceptor +): + async def intercept_stream_stream( + self, continuation, client_call_details, request_iterator + ): + return await continuation(self._attach(client_call_details), request_iterator) + + +def _async_bearer_interceptors( + token_provider: Callable[[], str], +) -> builtins.list[grpc.aio.ClientInterceptor]: + """One bearer interceptor per `grpc.aio` call category. + + See `_AsyncBearerAuthMixin` for why every category needs its own + instance rather than a single all-interfaces interceptor. + """ + return [ + _AsyncBearerUnaryUnaryInterceptor(token_provider), + _AsyncBearerUnaryStreamInterceptor(token_provider), + _AsyncBearerStreamUnaryInterceptor(token_provider), + _AsyncBearerStreamStreamInterceptor(token_provider), + ] + + +class AsyncSandboxSession: + """Async twin of `SandboxSession` bound to one sandbox id.""" + + def __init__(self, client: AsyncSandboxClient, sandbox: SandboxRef) -> None: + self._client = client + self.sandbox = sandbox + self._workspace = sandbox.workspace + + @property + def id(self) -> str: + return self.sandbox.id + + async def exec( + self, + command: Sequence[str], + *, + stream_output: bool = False, + workdir: str | None = None, + env: Mapping[str, str] | None = None, + stdin: bytes | None = None, + timeout_seconds: int | None = None, + ) -> ExecResult: + return await self._client.exec( + self.sandbox.id, + command, + stream_output=stream_output, + workdir=workdir, + env=env, + stdin=stdin, + timeout_seconds=timeout_seconds, + ) + + async def exec_python( + self, + function: Callable[..., object], + *, + args: Sequence[object] = (), + kwargs: Mapping[str, object] | None = None, + stream_output: bool = False, + workdir: str | None = None, + env: Mapping[str, str] | None = None, + timeout_seconds: int | None = None, + ) -> ExecResult: + return await self._client.exec_python( + self.sandbox.id, + function, + args=args, + kwargs=kwargs, + stream_output=stream_output, + workdir=workdir, + env=env, + timeout_seconds=timeout_seconds, + ) + + async def delete(self) -> bool: + return await self._client.delete(self.sandbox.name, workspace=self._workspace) + + +class AsyncSandboxClient: + """Async gRPC client for sandbox CRUD and command execution. + + Async twin of `SandboxClient` built on `grpc.aio`. Reuses the same + proto stubs, `TlsConfig`, bearer-auth plumbing, and on-disk gateway + resolution; the method surface mirrors the sync client with + `async`/`await` and `async with` support. + """ + + def __init__( + self, + endpoint: str, + *, + tls: TlsConfig | None = None, + bearer_token: str | Callable[[], str] | None = None, + timeout: float = 30.0, + cluster_name: str | None = None, + _bearer_close: Callable[[], None] | None = None, + ) -> None: + """Create an AsyncSandboxClient. + + Args mirror `SandboxClient.__init__`. `grpc.aio` accepts client + interceptors at channel-construction time, so the bearer + interceptor (when a token is supplied) is attached via the + channel's `interceptors=` argument rather than by wrapping an + existing channel. + """ + self._endpoint = endpoint + self._timeout = timeout + self._cluster_name = cluster_name + self._bearer_close = _bearer_close + provider = _normalize_bearer(bearer_token) + interceptors: list[grpc.aio.ClientInterceptor] = ( + _async_bearer_interceptors(provider) if provider is not None else [] + ) + if tls is None: + self._channel = grpc.aio.insecure_channel( + endpoint, interceptors=interceptors + ) + else: + # Build credentials from whatever subset of mTLS material the + # caller supplied. None for `root_certificates` makes gRPC use + # the system trust store, which is what we want for OIDC + # gateways behind a public CA. + credentials = grpc.ssl_channel_credentials( + root_certificates=(tls.ca_path.read_bytes() if tls.ca_path else None), + private_key=(tls.key_path.read_bytes() if tls.key_path else None), + certificate_chain=( + tls.cert_path.read_bytes() if tls.cert_path else None + ), + ) + self._channel = grpc.aio.secure_channel( + endpoint, credentials, interceptors=interceptors + ) + self._stub = openshell_pb2_grpc.OpenShellStub(self._channel) + + @classmethod + def from_active_cluster( + cls, + *, + cluster: str | None = None, + timeout: float = 30.0, + auto_refresh: bool = True, + write_back: bool = True, + insecure: bool = False, + ) -> AsyncSandboxClient: + """Construct an `AsyncSandboxClient` from the active gateway's state. + + Resolves the same on-disk gateway layout as + `SandboxClient.from_active_cluster` (via + `_resolve_gateway_connection`); see that method for the full + semantics of `cluster`, `timeout`, and the OIDC kwargs + (`auto_refresh`, `write_back`, `insecure`). Construction itself + performs no network I/O — the `grpc.aio` channel connects lazily + on the first RPC. + """ + conn = _resolve_gateway_connection( + cluster=cluster, + auto_refresh=auto_refresh, + write_back=write_back, + insecure=insecure, + ) + return cls( + conn.endpoint, + tls=conn.tls, + bearer_token=conn.bearer_token, + timeout=timeout, + cluster_name=conn.cluster_name, + _bearer_close=conn.bearer_close, + ) + + async def close(self) -> None: + """Release the gRPC channel and any bearer-auth resources. + + Idempotent. Mirrors `SandboxClient.close`, awaiting the + `grpc.aio` channel's async `close()`. If `from_active_cluster` + wired up an OIDC refresher for this client, the refresher's + underlying httpx.Client is closed here too. + """ + await self._channel.close() + if self._bearer_close is not None: + with contextlib.suppress(Exception): + self._bearer_close() + self._bearer_close = None + + async def __aenter__(self) -> AsyncSandboxClient: + return self + + async def __aexit__(self, *args: object) -> None: + await self.close() + + async def health(self) -> openshell_pb2.HealthResponse: + return await self._stub.Health( + openshell_pb2.HealthRequest(), timeout=self._timeout + ) + + async def create( + self, + *, + workspace: str, + spec: openshell_pb2.SandboxSpec | None = None, + name: str | None = None, + labels: Mapping[str, str] | None = None, + ) -> SandboxRef: + request_spec = spec if spec is not None else _default_spec() + response = await self._stub.CreateSandbox( + openshell_pb2.CreateSandboxRequest( + spec=request_spec, + name=name or "", + labels=dict(labels) if labels else {}, + workspace=workspace, + ), + timeout=self._timeout, + ) + sandbox_ref = _sandbox_ref(response.sandbox) + if sandbox_ref.id == "": + raise SandboxError("CreateSandbox returned empty sandbox id") + return sandbox_ref + + async def create_session( + self, + *, + workspace: str, + spec: openshell_pb2.SandboxSpec | None = None, + name: str | None = None, + labels: Mapping[str, str] | None = None, + ) -> AsyncSandboxSession: + return AsyncSandboxSession( + self, + await self.create(workspace=workspace, spec=spec, name=name, labels=labels), + ) + + async def get(self, sandbox_name: str, *, workspace: str) -> SandboxRef: + response = await self._stub.GetSandbox( + openshell_pb2.GetSandboxRequest(name=sandbox_name, workspace=workspace), + timeout=self._timeout, + ) + return _sandbox_ref(response.sandbox) + + async def get_session( + self, sandbox_name: str, *, workspace: str + ) -> AsyncSandboxSession: + return AsyncSandboxSession( + self, await self.get(sandbox_name, workspace=workspace) + ) + + async def list( + self, + *, + workspace: str, + limit: int = 100, + offset: int = 0, + label_selector: str | None = None, + ) -> builtins.list[SandboxRef]: + response = await self._stub.ListSandboxes( + openshell_pb2.ListSandboxesRequest( + workspace=workspace, + limit=limit, + offset=offset, + label_selector=label_selector or "", + ), + timeout=self._timeout, + ) + return [_sandbox_ref(item) for item in response.sandboxes] + + async def list_for_all_workspaces( + self, + *, + limit: int = 100, + offset: int = 0, + label_selector: str | None = None, + ) -> builtins.list[SandboxRef]: + response = await self._stub.ListSandboxes( + openshell_pb2.ListSandboxesRequest( + all_workspaces=True, + limit=limit, + offset=offset, + label_selector=label_selector or "", + ), + timeout=self._timeout, + ) + return [_sandbox_ref(item) for item in response.sandboxes] + + async def list_ids( + self, + *, + workspace: str, + limit: int = 100, + offset: int = 0, + label_selector: str | None = None, + ) -> builtins.list[str]: + return [ + item.id + for item in await self.list( + workspace=workspace, + limit=limit, + offset=offset, + label_selector=label_selector, + ) + ] + + async def list_ids_for_all_workspaces( + self, + *, + limit: int = 100, + offset: int = 0, + label_selector: str | None = None, + ) -> builtins.list[str]: + return [ + item.id + for item in await self.list_for_all_workspaces( + limit=limit, + offset=offset, + label_selector=label_selector, + ) + ] + + async def delete(self, sandbox_name: str, *, workspace: str) -> bool: + response = await self._stub.DeleteSandbox( + openshell_pb2.DeleteSandboxRequest(name=sandbox_name, workspace=workspace), + timeout=self._timeout, + ) + return bool(response.deleted) + + async def wait_deleted( + self, + sandbox_name: str, + *, + workspace: str, + timeout_seconds: float = 60.0, + ) -> None: + deadline = time.time() + timeout_seconds + while time.time() < deadline: + try: + await self.get(sandbox_name, workspace=workspace) + except grpc.aio.AioRpcError as exc: + if exc.code() == grpc.StatusCode.NOT_FOUND: + return + raise + await asyncio.sleep(1) + raise SandboxError(f"sandbox {sandbox_name} was not deleted within timeout") + + async def wait_ready( + self, + sandbox_name: str, + *, + workspace: str, + timeout_seconds: float = 300.0, + ) -> SandboxRef: + deadline = time.time() + timeout_seconds + while time.time() < deadline: + sandbox = await self.get(sandbox_name, workspace=workspace) + if sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_READY: + return sandbox + if sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_ERROR: + raise SandboxError(f"sandbox {sandbox_name} entered error phase") + await asyncio.sleep(1) + raise SandboxError(f"sandbox {sandbox_name} was not ready within timeout") + + async def exec_stream( + self, + sandbox_id: str, + command: Sequence[str], + *, + workdir: str | None = None, + env: Mapping[str, str] | None = None, + stdin: bytes | None = None, + timeout_seconds: int | None = None, + ) -> AsyncIterator[ExecChunk | ExecResult]: + if not command: + raise SandboxError("command must not be empty") + + request = openshell_pb2.ExecSandboxRequest( + sandbox_id=sandbox_id, + command=list(command), + workdir=workdir or "", + environment=dict(env or {}), + timeout_seconds=timeout_seconds or 0, + stdin=stdin or b"", + ) + # Use whichever is larger: the default client timeout or the command + # timeout plus headroom for SSH setup / teardown overhead. + grpc_deadline = self._timeout + if timeout_seconds and timeout_seconds + 10 > grpc_deadline: + grpc_deadline = timeout_seconds + 10 + stream = self._stub.ExecSandbox(request, timeout=grpc_deadline) + + stdout_parts: list[bytes] = [] + stderr_parts: list[bytes] = [] + exit_code: int | None = None + + async for event in stream: + payload = event.WhichOneof("payload") + if payload == "stdout": + data = bytes(event.stdout.data) + stdout_parts.append(data) + yield ExecChunk(stream="stdout", data=data) + elif payload == "stderr": + data = bytes(event.stderr.data) + stderr_parts.append(data) + yield ExecChunk(stream="stderr", data=data) + elif payload == "exit": + exit_code = int(event.exit.exit_code) + + if exit_code is None: + raise SandboxError("ExecSandbox stream ended without an exit event") + + yield ExecResult( + exit_code=exit_code, + stdout=b"".join(stdout_parts).decode("utf-8", errors="replace"), + stderr=b"".join(stderr_parts).decode("utf-8", errors="replace"), + ) + + async def exec( + self, + sandbox_id: str, + command: Sequence[str], + *, + stream_output: bool = False, + workdir: str | None = None, + env: Mapping[str, str] | None = None, + stdin: bytes | None = None, + timeout_seconds: int | None = None, + ) -> ExecResult: + result: ExecResult | None = None + async for item in self.exec_stream( + sandbox_id, + command, + workdir=workdir, + env=env, + stdin=stdin, + timeout_seconds=timeout_seconds, + ): + if stream_output and isinstance(item, ExecChunk): + if item.stream == "stdout": + sys.stdout.buffer.write(item.data) + sys.stdout.flush() + else: + sys.stderr.buffer.write(item.data) + sys.stderr.flush() + if isinstance(item, ExecResult): + result = item + if result is None: + raise SandboxError("ExecSandbox did not return a result") + return result + + async def exec_python( + self, + sandbox_id: str, + function: Callable[..., object], + *, + args: Sequence[object] = (), + kwargs: Mapping[str, object] | None = None, + stream_output: bool = False, + workdir: str | None = None, + env: Mapping[str, str] | None = None, + timeout_seconds: int | None = None, + ) -> ExecResult: + exec_env = dict(env or {}) + exec_env["OPENSHELL_PYFUNC_B64"] = _serialize_python_callable( + function, + args=args, + kwargs=kwargs, + ) + return await self.exec( + sandbox_id, + [_SANDBOX_PYTHON_BIN, "-c", _PYTHON_CLOUDPICKLE_BOOTSTRAP], + stream_output=stream_output, + workdir=workdir, + env=exec_env, + timeout_seconds=timeout_seconds, + ) + + +class AsyncInferenceRouteClient: + """Async twin of `InferenceRouteClient` for workspace inference routes.""" + + def __init__(self, channel: grpc.aio.Channel, *, timeout: float = 30.0) -> None: + self._stub = inference_pb2_grpc.InferenceStub(channel) + self._timeout = timeout + + @classmethod + def from_sandbox_client( + cls, client: AsyncSandboxClient + ) -> AsyncInferenceRouteClient: + return cls(client._channel, timeout=client._timeout) + + async def set_route( + self, + *, + workspace: str, + provider_name: str, + model_id: str, + no_verify: bool = False, + ) -> InferenceRouteConfig: + response = await self._stub.SetInferenceRoute( + inference_pb2.SetInferenceRouteRequest( + workspace=workspace, + provider_name=provider_name, + model_id=model_id, + no_verify=no_verify, + ), + timeout=self._timeout, + ) + return InferenceRouteConfig( + provider_name=response.provider_name, + model_id=response.model_id, + version=response.version, + ) + + async def get_route(self, *, workspace: str) -> InferenceRouteConfig: + response = await self._stub.GetInferenceRoute( + inference_pb2.GetInferenceRouteRequest(workspace=workspace), + timeout=self._timeout, + ) + return InferenceRouteConfig( + provider_name=response.provider_name, + model_id=response.model_id, + version=response.version, + ) + + async def delete_route( + self, + *, + workspace: str, + route_name: str = "", + ) -> bool: + response = await self._stub.DeleteInferenceRoute( + inference_pb2.DeleteInferenceRouteRequest( + workspace=workspace, + route_name=route_name, + ), + timeout=self._timeout, + ) + return response.deleted + + +class AsyncWorkspaceClient: + """Async twin of `WorkspaceClient` for workspace lifecycle operations.""" + + def __init__(self, channel: grpc.aio.Channel, *, timeout: float = 30.0) -> None: + self._stub = openshell_pb2_grpc.OpenShellStub(channel) + self._timeout = timeout + + @classmethod + def from_sandbox_client(cls, client: AsyncSandboxClient) -> AsyncWorkspaceClient: + return cls(client._channel, timeout=client._timeout) + + async def create( + self, + name: str, + *, + labels: Mapping[str, str] | None = None, + ) -> WorkspaceRef: + response = await self._stub.CreateWorkspace( + openshell_pb2.CreateWorkspaceRequest( + name=name, + labels=dict(labels) if labels else {}, + ), + timeout=self._timeout, + ) + return _workspace_ref(response.workspace) + + async def get(self, name: str) -> WorkspaceRef: + response = await self._stub.GetWorkspace( + openshell_pb2.GetWorkspaceRequest(name=name), + timeout=self._timeout, + ) + return _workspace_ref(response.workspace) + + async def list( + self, + *, + limit: int = 100, + offset: int = 0, + label_selector: str | None = None, + ) -> builtins.list[WorkspaceRef]: + response = await self._stub.ListWorkspaces( + openshell_pb2.ListWorkspacesRequest( + limit=limit, + offset=offset, + label_selector=label_selector or "", + ), + timeout=self._timeout, + ) + return [_workspace_ref(ws) for ws in response.workspaces] + + async def delete(self, name: str) -> bool: + response = await self._stub.DeleteWorkspace( + openshell_pb2.DeleteWorkspaceRequest(name=name), + timeout=self._timeout, + ) + return response.deleted + + +class AsyncSandbox: + """Async context-managed sandbox session bound to one sandbox id. + + Async twin of `Sandbox`: use it with `async with`. Constructor + arguments and OIDC kwargs mirror `Sandbox` exactly and forward to + `AsyncSandboxClient.from_active_cluster`. + """ + + def __init__( + self, + *, + workspace: str, + cluster: str | None = None, + sandbox: str | SandboxRef | None = None, + delete_on_exit: bool = True, + spec: openshell_pb2.SandboxSpec | None = None, + name: str | None = None, + labels: Mapping[str, str] | None = None, + timeout: float = 30.0, + ready_timeout_seconds: float = 120.0, + auto_refresh: bool = True, + write_back: bool = True, + insecure: bool = False, + ) -> None: + """Bind an AsyncSandbox context to the active gateway. + + OIDC kwargs (`auto_refresh`, `write_back`, `insecure`) forward + directly to `AsyncSandboxClient.from_active_cluster` and share + the same semantics as on `Sandbox`. + """ + self._workspace = workspace + self._cluster = cluster + self._sandbox_input = sandbox + self._delete_on_exit = delete_on_exit + self._spec = spec + self._name = name + # Copy so later caller mutation cannot change what gets sent on enter. + self._labels = dict(labels) if labels is not None else None + self._timeout = timeout + self._ready_timeout_seconds = ready_timeout_seconds + self._auto_refresh = auto_refresh + self._write_back = write_back + self._insecure = insecure + self._client: AsyncSandboxClient | None = None + self._session: AsyncSandboxSession | None = None + + @property + def id(self) -> str: + if self._session is None: + raise SandboxError("sandbox context has not been entered") + return self._session.id + + @property + def sandbox(self) -> SandboxRef: + if self._session is None: + raise SandboxError("sandbox context has not been entered") + return self._session.sandbox + + async def __aenter__(self) -> AsyncSandbox: + # Creation metadata cannot be applied when attaching to an existing + # sandbox; reject it before opening a connection. + if self._sandbox_input is not None and ( + self._name is not None or self._labels is not None + ): + raise SandboxError( + "name and labels cannot be set when attaching to an existing sandbox" + ) + + client = AsyncSandboxClient.from_active_cluster( + cluster=self._cluster, + timeout=self._timeout, + auto_refresh=self._auto_refresh, + write_back=self._write_back, + insecure=self._insecure, + ) + self._client = client + + if self._sandbox_input is None: + self._session = await client.create_session( + workspace=self._workspace, + spec=self._spec, + name=self._name, + labels=self._labels, + ) + elif isinstance(self._sandbox_input, SandboxRef): + self._session = AsyncSandboxSession(client, self._sandbox_input) + else: + self._session = await client.get_session( + self._sandbox_input, workspace=self._workspace + ) + + self._workspace = getattr(self._session, "_workspace", self._workspace) + + ready = await client.wait_ready( + self._session.sandbox.name, + workspace=self._workspace, + timeout_seconds=self._ready_timeout_seconds, + ) + self._session = AsyncSandboxSession(client, ready) + + return self + + async def __aexit__(self, *args: object) -> None: + try: + if ( + self._delete_on_exit + and self._session is not None + and self._client is not None + ): + try: + deleted = await self._session.delete() + if deleted: + await self._client.wait_deleted( + self._session.sandbox.name, + workspace=self._workspace, + ) + except grpc.aio.AioRpcError as exc: + if exc.code() != grpc.StatusCode.NOT_FOUND: + raise + finally: + if self._client is not None: + await self._client.close() + self._session = None + self._client = None + + async def exec( + self, + command: Sequence[str], + *, + stream_output: bool = False, + workdir: str | None = None, + env: Mapping[str, str] | None = None, + stdin: bytes | None = None, + timeout_seconds: int | None = None, + ) -> ExecResult: + if self._session is None: + raise SandboxError("sandbox context has not been entered") + return await self._session.exec( + command, + stream_output=stream_output, + workdir=workdir, + env=env, + stdin=stdin, + timeout_seconds=timeout_seconds, + ) + + async def exec_python( + self, + function: Callable[..., object], + *, + args: Sequence[object] = (), + kwargs: Mapping[str, object] | None = None, + stream_output: bool = False, + workdir: str | None = None, + env: Mapping[str, str] | None = None, + timeout_seconds: int | None = None, + ) -> ExecResult: + if self._session is None: + raise SandboxError("sandbox context has not been entered") + return await self._session.exec_python( + function, + args=args, + kwargs=kwargs, + stream_output=stream_output, + workdir=workdir, + env=env, + timeout_seconds=timeout_seconds, + )