From 6aae4ae77f86f1ef4027962eb222fd2fd17ca8f0 Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Mon, 2 Nov 2026 16:28:20 -0500 Subject: [PATCH 1/2] feat(python): add native async sandbox clients Signed-off-by: Max Dubrinsky --- .gitignore | 1 + docs/sdk/python.mdx | 41 +++- pyproject.toml | 2 +- python/README.md | 18 ++ python/grpc_typing_test.py | 23 ++ python/openshell/__init__.py | 14 ++ python/openshell/_async_runtime.py | 320 +++++++++++++++++++++++++ python/openshell/async_client_test.py | 205 ++++++++++++++++ python/openshell/async_routing_test.py | 107 +++++++++ python/openshell/providers.py | 20 +- python/openshell/sandbox.py | 209 +++++++++++----- python/wheel_verifier_test.py | 2 + tasks/scripts/generate_python_async.py | 265 ++++++++++++++++++++ tasks/scripts/generate_python_proto.py | 6 + tasks/scripts/verify-python-wheel.py | 2 + 15 files changed, 1160 insertions(+), 75 deletions(-) create mode 100644 python/README.md create mode 100644 python/openshell/_async_runtime.py create mode 100644 python/openshell/async_client_test.py create mode 100644 python/openshell/async_routing_test.py create mode 100644 tasks/scripts/generate_python_async.py diff --git a/.gitignore b/.gitignore index 4b110041a9..5f5361bdde 100644 --- a/.gitignore +++ b/.gitignore @@ -107,6 +107,7 @@ cython_debug/ # Generated Python protobuf stubs (keep package marker) python/openshell/_proto/* !python/openshell/_proto/__init__.py +python/openshell/aio.py # ============================================================================= # IDE / Editor diff --git a/docs/sdk/python.mdx b/docs/sdk/python.mdx index 165fd7ac8a..9cea0f5085 100644 --- a/docs/sdk/python.mdx +++ b/docs/sdk/python.mdx @@ -148,6 +148,45 @@ workspace for repeated exec, stop, start, and delete operations. List methods return lazy `Pager` instances; use `list_all()` only when you want to fetch the complete collection. +## Use the Async Client + +`AsyncSandboxClient` provides the same sandbox operations and arguments as +`SandboxClient` over `grpc.aio`. Construct and use it on one event loop: + +```python +from openshell import AsyncSandboxClient + +async def run(): + async with AsyncSandboxClient.from_active_cluster() as client: + sandbox = await client.get("worker", workspace="default") + result = await client.exec(sandbox.name, ["echo", "hello"], workspace="default") + print(result.stdout) + + async for page in client.list(workspace="default"): + for item in page.items: + print(item.name) +``` + +List methods return lazy `AsyncPager` objects without an `await`. Use +`await client.list_all(...)` to collect every page. The provider, sandbox-template, +and workspace clients also have `Async` counterparts. Authentication settings +match the synchronous client; credential refresh runs off the event loop. + +Both sandbox clients accept a caller-configured channel through `from_channel()`. +That factory borrows the channel: closing the client leaves it open. Clients +constructed from an endpoint or registered gateway own and close their channels. + +RPC methods accept `timeout` in seconds to override the client default. This is +separate from an exec command's `timeout_seconds`. Mutations that support gateway +deduplication accept `request_id`; reuse the same ID and payload when explicitly +retrying an uncertain result. Request IDs do not enable general automatic retries. Exec retains the gateway +owner-route fallback when the route fails before the first response; it never +replays a command after receiving output. + +Async task cancellation propagates to the active RPC. For early exit from +`exec_stream()`, use `contextlib.aclosing()` so the stream is cancelled promptly. +Gateway RPC failures raise `GatewayError`, preserving the gRPC status and details. + ## Call Generated RPCs The generated gRPC stubs include concrete request and response types for synchronous `grpc.Channel` and asynchronous `grpc.aio.Channel` connections. Install the dependency type stubs in your development environment for editor signatures and type checking: @@ -177,7 +216,7 @@ async def get_provider(channel: grpc.aio.Channel, name: str) -> pb.ProviderRespo ) ``` -Your application configures the channel's TLS and authentication, supplies RPC deadlines, and closes the channel. Create and use async channels on the same event loop. Cancellation and gRPC errors propagate directly. Generated stubs do not reuse `SandboxClient` authentication automatically; `SandboxClient` and `ProviderClient` remain synchronous. +Your application configures the channel's TLS and authentication, supplies RPC deadlines, and closes the channel. Create and use async channels on the same event loop. Cancellation and gRPC errors propagate directly. Generated stubs do not reuse SDK client authentication automatically. Use `AsyncSandboxClient` for the higher-level async methods and managed authentication. If you annotate a stored async stub explicitly, use a deferred annotation such as `"openshell_pb2_grpc.OpenShellAsyncStub"`. That type exists only in the `.pyi` file; construct `OpenShellStub` at runtime. The same convention applies to credential-driver stubs. diff --git a/pyproject.toml b/pyproject.toml index 857a69b450..ba5c0578ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,7 +78,7 @@ git_describe_command = ["git", "describe", "--dirty", "--tags", "--long", "--mat [tool.ruff] target-version = "py311" line-length = 88 -exclude = ["python/openshell/_proto"] +exclude = ["python/openshell/_proto", "python/openshell/aio.py"] [tool.ruff.lint] select = [ diff --git a/python/README.md b/python/README.md new file mode 100644 index 0000000000..11228d8713 --- /dev/null +++ b/python/README.md @@ -0,0 +1,18 @@ +# Python SDK implementation + +The synchronous clients in `openshell/sandbox.py` and `openshell/providers.py` +own request construction and response conversion. `tasks/scripts/generate_python_async.py` +emits concrete async counterparts into `openshell/aio.py`; the handwritten +`openshell/_async_runtime.py` owns asyncio transport, authentication, pagination, +and resource cleanup. Generated files are build outputs and are not committed. + +Run `mise run python:proto` after changing client methods or protobufs. The same +generation runs when building release wheels. The async transform recognizes +stub calls, lazy pager fetches, stream iteration, and polling sleeps. Add new I/O +shapes explicitly to the generator and test their runtime and consumer typing +contracts. Keep blocking operations out of generated async method bodies. + +Validate changes with Python formatting, lint, and type checks plus the affected +tests. `python/openshell/async_client_test.py` exercises real loopback gRPC calls; +`python/grpc_typing_test.py` checks consumer signatures. The wheel verifier requires +both generated bindings and the async client module. diff --git a/python/grpc_typing_test.py b/python/grpc_typing_test.py index cfb4342dc7..ec82bec4db 100644 --- a/python/grpc_typing_test.py +++ b/python/grpc_typing_test.py @@ -79,3 +79,26 @@ def _check_types(path: Path) -> subprocess.CompletedProcess[str]: text=True, check=False, ) + + +def test_async_sdk_consumer_types(tmp_path: Path) -> None: + source = """ +from typing import assert_type +from openshell import AsyncSandboxClient, AsyncPager, SandboxRef, ExecResult, DeletionResult + +async def check(client: AsyncSandboxClient) -> None: + assert_type(await client.get("worker", workspace="team"), SandboxRef) + assert_type(await client.create(workspace="team"), SandboxRef) + assert_type(client.list(workspace="team"), AsyncPager[SandboxRef]) + assert_type(await client.list_all(workspace="team"), list[SandboxRef]) + assert_type(await client.exec("worker", ["echo"], workspace="team"), ExecResult) + assert_type(await client.delete("worker", workspace="team"), DeletionResult) +""" + path = tmp_path / "sdk_consumer.py" + path.write_text(source) + result = _check_types(path) + assert result.returncode == 0, result.stdout + result.stderr + path.write_text(source.replace('workspace="team"', "workspace=123")) + result = _check_types(path) + assert result.returncode != 0 + assert "invalid-argument-type" in result.stdout diff --git a/python/openshell/__init__.py b/python/openshell/__init__.py index c99c3b48e0..fdfadecbb3 100644 --- a/python/openshell/__init__.py +++ b/python/openshell/__init__.py @@ -5,6 +5,14 @@ from __future__ import annotations +from ._async_runtime import AsyncPager +from .aio import ( + AsyncProviderClient, + AsyncSandboxClient, + AsyncSandboxSession, + AsyncSandboxTemplateClient, + AsyncWorkspaceClient, +) from .errors import ErrorInfo, FieldViolation, GatewayError, from_grpc_error from .mutations import DeletionOutcome, DeletionResult from .providers import ProviderClient, ProviderCredentialValue @@ -38,6 +46,12 @@ __version__ = "0.0.0" __all__ = [ + "AsyncPager", + "AsyncProviderClient", + "AsyncSandboxClient", + "AsyncSandboxSession", + "AsyncSandboxTemplateClient", + "AsyncWorkspaceClient", "ClientCredentialsAuth", "DeletionOutcome", "DeletionResult", diff --git a/python/openshell/_async_runtime.py b/python/openshell/_async_runtime.py new file mode 100644 index 0000000000..30651b96c4 --- /dev/null +++ b/python/openshell/_async_runtime.py @@ -0,0 +1,320 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native asyncio transport and iteration for the generated async clients.""" + +from __future__ import annotations + +import asyncio +import contextlib +import time +from typing import TYPE_CHECKING, Generic, Self, TypeVar + +import grpc +import grpc.aio + +from ._proto import openshell_pb2, openshell_pb2_grpc +from .errors import from_grpc_error +from .sandbox import ( + _OWNER_REPLICA_HEADER, + _OWNER_REQUEST_HEADER, + _ROUTE_REPLICA_HEADER, + ClientCredentialsAuth, + Page, + Pager, + SandboxError, + TlsConfig, + _get_sandbox_request, + _is_local_grpc_endpoint, + _normalize_bearer, + _OwnerCache, + _retry_unrouted, +) + +T = TypeVar("T") + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + +async def _await_rpc(call: Awaitable[T]) -> T: + try: + return await call + except grpc.RpcError as error: + raise from_grpc_error(error) from error + + +class AsyncPager(Generic[T]): + """Lazy, single-pass async pages with the same token guards as ``Pager``. + + A pager must have only one consumer. Cancellation leaves the current token + available for an explicit retry; no page is prefetched or retried implicitly. + """ + + def __init__( + self, fetch: Callable[[str], Awaitable[Page[T]]], page_token: str = "" + ) -> None: + self._fetch = fetch + # Reuse the sync pager's bounded token accounting, not its I/O. + self._state: Pager[T] = Pager(lambda _: Page([], ""), page_token) + self._busy = False + + def __aiter__(self) -> AsyncPager[T]: + return self + + async def __anext__(self) -> Page[T]: + if self._busy: + raise RuntimeError("an AsyncPager cannot have concurrent consumers") + state = self._state + if state._page_token is None: + raise StopAsyncIteration + token = state._page_token + token_bytes = state._validate_page_token_budget(token) + self._busy = True + try: + page = await self._fetch(token) + if token: + state._consumed_page_tokens.add(token) + state._consumed_page_token_bytes += token_bytes + if ( + page.next_page_token + and page.next_page_token in state._consumed_page_tokens + ): + raise SandboxError("pager received a repeated continuation token") + state._page_token = page.next_page_token or None + return page + finally: + self._busy = False + + async def all(self) -> list[T]: + """Collect every remaining page into memory.""" + return [item async for page in self for item in page.items] + + +class _AsyncBearer: + """Run the existing locked refresh providers off the event loop. + + Keep worker tasks alive across caller cancellation so close can wait before + disposing an owned refresher. Providers retain their existing single-flight + refresh and cross-process rotated-token persistence semantics. + """ + + def __init__(self, provider: Callable[[], str]) -> None: + self._provider = provider + self._pending: set[asyncio.Task[str]] = set() + + async def details(self, details: grpc.aio.ClientCallDetails): + start = time.monotonic() + task = asyncio.create_task(asyncio.to_thread(self._provider)) + self._pending.add(task) + task.add_done_callback(self._finished) + try: + token = await asyncio.wait_for(asyncio.shield(task), details.timeout) + except TimeoutError as error: + raise grpc.aio.AioRpcError( + grpc.StatusCode.DEADLINE_EXCEEDED, + grpc.aio.Metadata(), + grpc.aio.Metadata(), + "deadline exceeded while obtaining bearer credentials", + ) from error + timeout = details.timeout + if timeout is not None: + timeout = max(0.0, timeout - (time.monotonic() - start)) + return grpc.aio.ClientCallDetails( + details.method, + timeout, + grpc.aio.Metadata( + *tuple(details.metadata or ()), ("authorization", f"Bearer {token}") + ), + details.credentials, + details.wait_for_ready, + ) + + def _finished(self, task: asyncio.Task[str]) -> None: + self._pending.discard(task) + if not task.cancelled(): + task.exception() # Retrieve failures when the requesting task was cancelled. + + async def drain(self) -> None: + if self._pending: + await asyncio.gather(*self._pending, return_exceptions=True) + + +class _UnaryAuth(grpc.aio.UnaryUnaryClientInterceptor): + def __init__(self, auth: _AsyncBearer) -> None: + self._auth = auth + + async def intercept_unary_unary(self, continuation, client_call_details, request): + return await continuation( + await self._auth.details(client_call_details), request + ) + + +class _StreamAuth(grpc.aio.UnaryStreamClientInterceptor): + def __init__(self, auth: _AsyncBearer) -> None: + self._auth = auth + + async def intercept_unary_stream(self, continuation, client_call_details, request): + return await continuation( + await self._auth.details(client_call_details), request + ) + + +class _BidiAuth(grpc.aio.StreamStreamClientInterceptor): + def __init__(self, auth: _AsyncBearer) -> None: + self._auth = auth + + async def intercept_stream_stream( + self, continuation, client_call_details, request_iterator + ): + return await continuation( + await self._auth.details(client_call_details), request_iterator + ) + + +class AsyncTransport: + """Transport owned by AsyncSandboxClient; construct on its event loop.""" + + def __init__( + self, + endpoint: str, + *, + tls: TlsConfig | None = None, + bearer_token: str | Callable[[], str] | None = None, + client_credentials: ClientCredentialsAuth | None = None, + timeout: float = 30.0, + cluster_name: str | None = None, + _bearer_close: Callable[[], None] | None = None, + ) -> None: + if bearer_token is not None and client_credentials is not None: + raise SandboxError( + "bearer_token and client_credentials are mutually exclusive" + ) + if ( + client_credentials is not None + and tls is None + and not _is_local_grpc_endpoint(endpoint) + ): + raise SandboxError( + "OAuth client credentials require TLS for non-loopback gateway endpoints" + ) + self._timeout = timeout + self._endpoint = endpoint + self._cluster_name = cluster_name + self._bearer_close = _bearer_close + self._owns_channel = True + self._close_task: asyncio.Task[None] | None = None + provider = _normalize_bearer(client_credentials or bearer_token) + self._auth = _AsyncBearer(provider) if provider is not None else None + # aio registers multi-interface interceptors under only their first + # matching interface, so use a separate object for each RPC shape. + interceptors = ( + [_UnaryAuth(self._auth), _StreamAuth(self._auth), _BidiAuth(self._auth)] + if self._auth is not None + else [] + ) + if tls is None: + self._channel = grpc.aio.insecure_channel( + endpoint, interceptors=interceptors + ) + else: + 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) + self._owners = _OwnerCache() + + @classmethod + def from_channel(cls, channel: grpc.aio.Channel, *, timeout: float = 30.0) -> Self: + """Borrow a caller-configured channel; closing the client leaves it open.""" + client = cls.__new__(cls) + client._channel = channel + client._stub = openshell_pb2_grpc.OpenShellStub(channel) + client._timeout = timeout + client._owns_channel = False + client._bearer_close = None + client._auth = None + client._close_task = None + client._owners = _OwnerCache() + return client + + async def _owner_of(self, workspace: str, sandbox: str, timeout: float) -> str: + owner = self._owners.get(workspace, sandbox) + if owner is not None: + return owner + call = self._stub.GetSandbox( + _get_sandbox_request(sandbox, workspace), + timeout=timeout, + metadata=((_OWNER_REQUEST_HEADER, "1"),), + ) + try: + await call + metadata = await call.initial_metadata() + except grpc.RpcError: + return "" + value = dict(tuple(metadata or ())).get(_OWNER_REPLICA_HEADER) + owner = value if isinstance(value, str) else "" + self._owners.put(workspace, sandbox, owner) + return owner + + async def _open_exec( + self, request: openshell_pb2.ExecSandboxRequest, workspace: str, timeout: float + ): + deadline = time.monotonic() + timeout + owner = await self._owner_of(workspace, request.sandbox, timeout) + if owner: + events = self._stub.ExecSandbox( + request, + timeout=max(0.0, deadline - time.monotonic()), + metadata=((_ROUTE_REPLICA_HEADER, owner),), + ) + try: + try: + first = await anext(events.__aiter__(), None) + except grpc.RpcError as error: + if not _retry_unrouted(error): + raise + self._owners.forget(workspace, request.sandbox) + else: + if first is not None: + yield first + async for event in events: + yield event + return + finally: + events.cancel() + events = self._stub.ExecSandbox( + request, timeout=max(0.0, deadline - time.monotonic()) + ) + try: + async for event in events: + yield event + finally: + events.cancel() + + async def _close(self) -> None: + if self._owns_channel: + await self._channel.close() + if self._auth is not None: + await self._auth.drain() + if self._bearer_close is not None: + with contextlib.suppress(Exception): + await asyncio.to_thread(self._bearer_close) + self._bearer_close = None + + async def close(self) -> None: + """Close owned resources once, even if the caller is cancelled.""" + if self._close_task is None: + self._close_task = asyncio.create_task(self._close()) + await asyncio.shield(self._close_task) + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *args: object) -> None: + await self.close() diff --git a/python/openshell/async_client_test.py b/python/openshell/async_client_test.py new file mode 100644 index 0000000000..f501f16e06 --- /dev/null +++ b/python/openshell/async_client_test.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Consumer contracts exercised over a real loopback aio transport.""" + +import asyncio +import inspect +import threading +from contextlib import asynccontextmanager + +import grpc +import grpc.aio +import pytest + +from openshell import ( + AsyncPager, + AsyncSandboxClient, + ExecChunk, + GatewayError, + Page, + SandboxError, +) +from openshell._proto import datamodel_pb2 as dm +from openshell._proto import openshell_pb2 as pb +from openshell._proto import openshell_pb2_grpc as rpc + + +class Gateway(rpc.OpenShellServicer): + def __init__(self): + self.requests = [] + self.started = asyncio.Event() + self.cancelled = asyncio.Event() + + async def GetSandbox(self, request, context): + self.requests.append((request, context.invocation_metadata())) + if request.name == "missing": + await context.abort(grpc.StatusCode.NOT_FOUND, "missing sandbox") + if request.name == "blocked": + self.started.set() + try: + await asyncio.Future() + finally: + self.cancelled.set() + return pb.SandboxResponse( + sandbox=pb.Sandbox( + metadata=dm.ObjectMeta( + id="identity", + name=request.name, + workspace=request.workspace_scope.workspace, + ), + status=pb.SandboxStatus(phase=pb.SANDBOX_PHASE_READY, exit_code=0), + ) + ) + + async def CreateSandbox(self, request, context): + self.requests.append((request, context.invocation_metadata())) + return pb.SandboxResponse( + sandbox=pb.Sandbox(metadata=dm.ObjectMeta(id="created", name=request.name)) + ) + + async def ListSandboxes(self, request, context): + self.requests.append((request, context.invocation_metadata())) + return pb.ListSandboxesResponse( + sandboxes=[ + pb.Sandbox(metadata=dm.ObjectMeta(name=request.page_token or "first")) + ], + next_page_token="second" if not request.page_token else "", + ) + + async def ExecSandbox(self, request, context): + self.requests.append((request, context.invocation_metadata())) + event = pb.ExecSandboxEvent() + event.stdout.data = b"hello" + yield event + if request.command == ["blocked"]: + self.started.set() + try: + await asyncio.Future() + finally: + self.cancelled.set() + event = pb.ExecSandboxEvent() + event.exit.exit_code = 0 + yield event + + +@asynccontextmanager +async def gateway(): + service = Gateway() + server = grpc.aio.server() + rpc.add_OpenShellServicer_to_server(service, server) + port = server.add_insecure_port("127.0.0.1:0") + await server.start() + try: + yield (service, f"127.0.0.1:{port}") + finally: + await server.stop(0) + + +async def test_async_lifecycle_pagination_exec_and_auth(): + async with ( + gateway() as (service, target), + AsyncSandboxClient(target, bearer_token="test-token") as client, + ): + sandbox = await client.get("worker", workspace="team") + assert sandbox.workspace == "team" + assert sandbox.status.exit_code == 0 + created = await client.create( + name="new", workspace="team", request_id="request-1" + ) + assert created.id == "created" + assert service.requests[-1][0].request_id == "request-1" + pages = client.list(workspace="team", label_selector="owner=helix") + assert isinstance(pages, AsyncPager) + assert [item.name for item in await pages.all()] == ["first", "second"] + result = await client.exec("worker", ["echo"], workspace="team") + assert result.stdout == "hello" + assert result.exit_code == 0 + for _, metadata in service.requests: + assert ("authorization", "Bearer test-token") in metadata + + +async def test_errors_deadlines_and_cancellation(): + async with gateway() as (service, target), AsyncSandboxClient(target) as client: + with pytest.raises(GatewayError) as failure: + await client.get("missing", workspace="team") + assert failure.value.code() == grpc.StatusCode.NOT_FOUND + with pytest.raises(GatewayError) as failure: + await client.get("blocked", workspace="team", timeout=0.05) + assert failure.value.code() == grpc.StatusCode.DEADLINE_EXCEEDED + await asyncio.wait_for(service.cancelled.wait(), 2) + service.started.clear() + service.cancelled.clear() + task = asyncio.create_task(client.get("blocked", workspace="team")) + await asyncio.wait_for(service.started.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(service.cancelled.wait(), 2) + assert (await client.get("healthy", workspace="team")).name == "healthy" + + +async def test_stream_close_cancels_rpc(): + async with ( + gateway() as (service, target), + AsyncSandboxClient(target, bearer_token="stream-token") as client, + ): + stream = client.exec_stream("worker", ["blocked"], workspace="team") + chunk = await anext(stream) + assert isinstance(chunk, ExecChunk) + assert chunk.data == b"hello" + await asyncio.wait_for(service.started.wait(), 2) + await stream.aclose() + await asyncio.wait_for(service.cancelled.wait(), 2) + + +async def test_borrowed_channel_remains_open(): + async with gateway() as (_, target), grpc.aio.insecure_channel(target) as channel: + async with AsyncSandboxClient.from_channel(channel) as client: + assert (await client.get("worker", workspace="team")).name == "worker" + response = await rpc.OpenShellStub(channel).GetSandbox( + pb.GetSandboxRequest(name="after-close") + ) + assert response.sandbox.metadata.name == "after-close" + + +async def test_blocking_token_provider_does_not_block_loop(): + entered = threading.Event() + release = threading.Event() + + def token(): + entered.set() + assert release.wait(5) + return "refreshed" + + async with ( + gateway() as (service, target), + AsyncSandboxClient(target, bearer_token=token) as client, + ): + pending = asyncio.create_task(client.get("worker", workspace="team")) + try: + assert await asyncio.to_thread(entered.wait, 2) + await asyncio.wait_for(asyncio.sleep(0), 0.5) + finally: + release.set() + await pending + assert ("authorization", "Bearer refreshed") in service.requests[0][1] + + +async def test_pager_rejects_repeated_tokens(): + async def fetch(token): + return Page([token], "again") + + pager = AsyncPager(fetch) + await anext(pager) + with pytest.raises(SandboxError, match="repeated"): + await anext(pager) + + +def test_async_methods_are_real_coroutines(): + from unittest.mock import AsyncMock, create_autospec + + assert inspect.iscoroutinefunction(AsyncSandboxClient.create) + assert not inspect.iscoroutinefunction(AsyncSandboxClient.list) + client = create_autospec(AsyncSandboxClient, instance=True) + assert isinstance(client.create, AsyncMock) diff --git a/python/openshell/async_routing_test.py b/python/openshell/async_routing_test.py new file mode 100644 index 0000000000..8858bb845c --- /dev/null +++ b/python/openshell/async_routing_test.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native async exec preserves replica affinity and the pre-response fallback.""" + +from contextlib import asynccontextmanager + +import grpc +import grpc.aio +import pytest + +from openshell import AsyncSandboxClient, ExecChunk, GatewayError +from openshell._proto import openshell_pb2 as pb +from openshell._proto import openshell_pb2_grpc as rpc + + +class Gateway(rpc.OpenShellServicer): + def __init__(self): + self.owner = "" + self.reject_route = None + self.fail_after_output = False + self.exec_deadlines = [] + self.requests = [] + + async def GetSandbox(self, request, context): + self.requests.append((request, context.invocation_metadata())) + if self.owner: + await context.send_initial_metadata((("x-openshell-owner", self.owner),)) + return pb.SandboxResponse() + + async def ExecSandbox(self, request, context): + self.requests.append((request, context.invocation_metadata())) + self.exec_deadlines.append(context.time_remaining()) + if ( + dict(context.invocation_metadata()).get("x-openshell-replica") + and self.reject_route is not None + ): + await context.abort(grpc.StatusCode.UNAVAILABLE, self.reject_route) + event = pb.ExecSandboxEvent() + event.stdout.data = b"hello" + yield event + if self.fail_after_output: + await context.abort(grpc.StatusCode.UNAVAILABLE, "no healthy upstream") + event = pb.ExecSandboxEvent() + event.exit.exit_code = 0 + yield event + + +@asynccontextmanager +async def gateway(): + service = Gateway() + server = grpc.aio.server() + rpc.add_OpenShellServicer_to_server(service, server) + port = server.add_insecure_port("127.0.0.1:0") + await server.start() + try: + yield service, f"127.0.0.1:{port}" + finally: + await server.stop(0) + + +async def test_async_exec_owner_cache_and_narrow_unrouted_retry(): + async with gateway() as (service, target), AsyncSandboxClient(target) as client: + service.owner = "openshell-1" + for _ in range(2): + assert ( + await client.exec("worker", ["echo"], workspace="team") + ).stdout == "hello" + lookups = [ + request + for request, _ in service.requests + if isinstance(request, pb.GetSandboxRequest) + ] + assert len(lookups) == 1 + service.reject_route = "no healthy upstream" + assert ( + await client.exec("worker", ["echo"], workspace="team") + ).stdout == "hello" + routes = [ + dict(metadata).get("x-openshell-replica", "") + for request, metadata in service.requests + if isinstance(request, pb.ExecSandboxRequest) + ] + assert routes == ["openshell-1", "openshell-1", "openshell-1", ""] + service.reject_route = "gateway denied exec" + with pytest.raises(GatewayError): + await client.exec("worker", ["echo"], workspace="team") + assert dict(service.requests[-1][1])["x-openshell-replica"] == "openshell-1" + + +async def test_async_exec_never_replays_after_output_and_preserves_deadlines(): + async with gateway() as (service, target), AsyncSandboxClient(target) as client: + service.owner = "openshell-1" + await client.exec("worker", ["echo"], workspace="team", timeout_seconds=60) + assert 60 < service.exec_deadlines[-1] <= 70 + async for _ in client.exec_stream( + "worker", ["echo"], workspace="team", timeout_seconds=60, timeout=0.5 + ): + pass + assert 0 < service.exec_deadlines[-1] <= 0.6 + service.fail_after_output = True + before = len(service.exec_deadlines) + stream = client.exec_stream("worker", ["echo"], workspace="team") + assert isinstance(await anext(stream), ExecChunk) + with pytest.raises(GatewayError): + await anext(stream) + assert len(service.exec_deadlines) == before + 1 diff --git a/python/openshell/providers.py b/python/openshell/providers.py index f3312f5407..a868ba8acd 100644 --- a/python/openshell/providers.py +++ b/python/openshell/providers.py @@ -57,6 +57,7 @@ def get_credentials( *, workspace: str, minimum_remaining_lifetime: timedelta | None = None, + timeout: float | None = None, ) -> dict[str, ProviderCredentialValue]: """Export exactly the selected runtime keys, or raise without a result. @@ -69,21 +70,20 @@ def get_credentials( raise ValueError("an explicit canonical workspace is required") if not name or name.strip() != name or len(name.encode("utf-8")) > 253: raise ValueError("a canonical provider name is required") - # Copy the selection so caller mutation cannot change response validation. if isinstance(credential_keys, str): raise ValueError("select between 1 and 32 unique runtime environment keys") keys = list(credential_keys) if ( not 1 <= len(keys) <= 32 or any( - len(key) > 256 or re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key) is None + len(key) > 256 or re.fullmatch("[A-Za-z_][A-Za-z0-9_]*", key) is None for key in keys ) or len(set(keys)) != len(keys) ): raise ValueError("select between 1 and 32 unique runtime environment keys") - if minimum_remaining_lifetime is not None and not ( - timedelta(0) <= minimum_remaining_lifetime <= timedelta(days=1) + if minimum_remaining_lifetime is not None and ( + not timedelta(0) <= minimum_remaining_lifetime <= timedelta(days=1) ): raise ValueError( "minimum_remaining_lifetime must be between zero and 24 hours" @@ -95,17 +95,17 @@ def get_credentials( ) if minimum_remaining_lifetime is not None: request.minimum_remaining_lifetime.FromTimedelta(minimum_remaining_lifetime) - response = self._stub.GetProviderCredentials(request, timeout=self._timeout) + response = self._stub.GetProviderCredentials( + request, timeout=self._timeout if timeout is None else timeout + ) if set(response.credentials) != set(keys): raise ValueError("gateway returned an unexpected credential selection") return { key: ProviderCredentialValue( value=credential.value, - expiration_time=( - credential.expiration_time.ToDatetime(tzinfo=UTC) - if credential.HasField("expiration_time") - else None - ), + expiration_time=credential.expiration_time.ToDatetime(tzinfo=UTC) + if credential.HasField("expiration_time") + else None, ) for key, credential in response.credentials.items() } diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index c02a33cdbf..092ea2506c 100644 --- a/python/openshell/sandbox.py +++ b/python/openshell/sandbox.py @@ -187,7 +187,7 @@ class _ClientCallDetails(_ClientCallDetailsBase, grpc.ClientCallDetails): if TYPE_CHECKING: import builtins - from collections.abc import Callable, Iterator, Mapping, Sequence + from collections.abc import Callable, Generator, Iterator, Mapping, Sequence from .providers import ProviderClient @@ -760,6 +760,7 @@ def __init__( self._timeout = timeout self._cluster_name = cluster_name self._bearer_close = _bearer_close + self._owns_channel = True if tls is None: self._channel = grpc.insecure_channel(endpoint) else: @@ -892,6 +893,20 @@ def from_active_cluster( _bearer_close=bearer_close, ) + @classmethod + def from_channel( + cls, channel: grpc.Channel, *, timeout: float = 30.0 + ) -> SandboxClient: + """Borrow a caller-configured channel without taking ownership of it.""" + client = cls.__new__(cls) + client._channel = channel + client._stub = openshell_pb2_grpc.OpenShellStub(_error_mapping_channel(channel)) + client._timeout = timeout + client._owns_channel = False + client._owners = _OwnerCache() + client._bearer_close = None + return client + def close(self) -> None: """Release the gRPC channel and any bearer-auth resources. @@ -900,7 +915,8 @@ def close(self) -> None: closed here too — otherwise long-lived services that churn clients would leak sockets / file descriptors until GC. """ - self._channel.close() + if self._owns_channel: + self._channel.close() if self._bearer_close is not None: with contextlib.suppress(Exception): self._bearer_close() @@ -912,8 +928,11 @@ def __enter__(self) -> SandboxClient: def __exit__(self, *args: object) -> None: self.close() - def health(self) -> openshell_pb2.HealthResponse: - return self._stub.Health(openshell_pb2.HealthRequest(), timeout=self._timeout) + def health(self, *, timeout: float | None = None) -> openshell_pb2.HealthResponse: + return self._stub.Health( + openshell_pb2.HealthRequest(), + timeout=self._timeout if timeout is None else timeout, + ) def create( self, @@ -923,6 +942,8 @@ def create( name: str | None = None, labels: Mapping[str, str] | None = None, service_exposures: Sequence[ServiceExposure] | None = None, + timeout: float | None = None, + request_id: str | None = None, ) -> SandboxRef: request_spec = spec if spec is not None else _default_spec() response = self._stub.CreateSandbox( @@ -932,8 +953,9 @@ def create( name=name or "", labels=dict(labels) if labels else {}, service_exposures=_service_exposure_messages(service_exposures), + request_id=request_id or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) sandbox_ref = _sandbox_ref(response.sandbox, response.service_urls) if sandbox_ref.id == "": @@ -949,6 +971,8 @@ def create_from_template( name: str | None = None, labels: Mapping[str, str] | None = None, service_exposures: Sequence[ServiceExposure] | None = None, + timeout: float | None = None, + request_id: str | None = None, ) -> SandboxRef: if not workload_template.strip(): raise SandboxError("workload_template is required") @@ -961,8 +985,9 @@ def create_from_template( labels=dict(labels) if labels else {}, workload_template=workload_template, service_exposures=_service_exposure_messages(service_exposures), + request_id=request_id or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) sandbox_ref = _sandbox_ref(response.sandbox, response.service_urls) if sandbox_ref.id == "": @@ -1020,10 +1045,12 @@ def providers(self) -> ProviderClient: return ProviderClient(self._channel, timeout=self._timeout) - def get(self, name: str, *, workspace: str) -> SandboxRef: + def get( + self, name: str, *, workspace: str, timeout: float | None = None + ) -> SandboxRef: response = self._stub.GetSandbox( _get_sandbox_request(name, workspace), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return _sandbox_ref(response.sandbox) @@ -1037,6 +1064,7 @@ def list( page_size: int = 100, page_token: str = "", label_selector: str | None = None, + timeout: float | None = None, ) -> Pager[SandboxRef]: def fetch(token: str) -> Page[SandboxRef]: response = self._stub.ListSandboxes( @@ -1046,7 +1074,7 @@ def fetch(token: str) -> Page[SandboxRef]: page_token=token, label_selector=label_selector or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return Page( items=[_sandbox_ref(item) for item in response.sandboxes], @@ -1076,6 +1104,7 @@ def list_for_all_workspaces( page_size: int = 100, page_token: str = "", label_selector: str | None = None, + timeout: float | None = None, ) -> Pager[SandboxRef]: def fetch(token: str) -> Page[SandboxRef]: response = self._stub.ListSandboxes( @@ -1085,7 +1114,7 @@ def fetch(token: str) -> Page[SandboxRef]: page_token=token, label_selector=label_selector or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return Page( items=[_sandbox_ref(item) for item in response.sandboxes], @@ -1138,41 +1167,66 @@ def list_ids_for_all_workspaces( ] def delete( - self, name: str, *, workspace: str, allow_missing: bool = False + self, + name: str, + *, + workspace: str, + allow_missing: bool = False, + timeout: float | None = None, + request_id: str | None = None, ) -> DeletionResult: response = self._stub.DeleteSandbox( openshell_pb2.DeleteSandboxRequest( workspace_scope=_workspace_scope(workspace), name=name, allow_missing=allow_missing, + request_id=request_id or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return DeletionResult( DeletionOutcome(response.outcome), response.sandbox_id or None ) - def stop(self, name: str, *, workspace: str) -> SandboxRef: + def stop( + self, + name: str, + *, + workspace: str, + timeout: float | None = None, + request_id: str | None = None, + ) -> SandboxRef: response = self._stub.StopSandbox( openshell_pb2.StopSandboxRequest( workspace_scope=_workspace_scope(workspace), name=name, + request_id=request_id or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return _sandbox_ref(response.sandbox) - def start(self, name: str, *, workspace: str) -> SandboxRef: + def start( + self, + name: str, + *, + workspace: str, + timeout: float | None = None, + request_id: str | None = None, + ) -> SandboxRef: response = self._stub.StartSandbox( openshell_pb2.StartSandboxRequest( workspace_scope=_workspace_scope(workspace), name=name, + request_id=request_id or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return _sandbox_ref(response.sandbox) - def suspend(self, name: str, *, workspace: str) -> SandboxRef: + def suspend( + self, name: str, *, workspace: str, timeout: float | None = None + ) -> SandboxRef: """Suspend a ready sandbox, retaining its memory state. Requires a compute driver and isolation backend that support suspend @@ -1180,10 +1234,9 @@ def suspend(self, name: str, *, workspace: str) -> SandboxRef: """ response = self._stub.SuspendSandbox( openshell_pb2.SuspendSandboxRequest( - workspace_scope=_workspace_scope(workspace), - name=name, + workspace_scope=_workspace_scope(workspace), name=name ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return _sandbox_ref(response.sandbox) @@ -1193,6 +1246,7 @@ def resume( *, workspace: str, restart_if_restore_fails: bool = False, + timeout: float | None = None, ) -> SandboxRef: """Resume a suspended sandbox in a new runtime generation.""" response = self._stub.ResumeSandbox( @@ -1201,7 +1255,7 @@ def resume( name=name, restart_if_restore_fails=restart_if_restore_fails, ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return _sandbox_ref(response.sandbox) @@ -1312,10 +1366,11 @@ def exec_stream( stdin: bytes | None = None, timeout_seconds: int | None = None, no_login_shell: bool = False, + timeout: float | None = None, + request_id: str | None = None, ) -> Iterator[ExecChunk | ExecResult]: if not command: raise SandboxError("command must not be empty") - request = openshell_pb2.ExecSandboxRequest( workspace_scope=_workspace_scope(workspace), sandbox=sandbox, @@ -1324,36 +1379,34 @@ def exec_stream( environment=dict(env or {}), stdin=stdin or b"", no_login_shell=no_login_shell, + request_id=request_id or "", ) if timeout_seconds: request.execution_timeout.seconds = timeout_seconds - # 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 = self._timeout if timeout is None else timeout + if timeout is None and timeout_seconds and timeout_seconds + 10 > grpc_deadline: grpc_deadline = timeout_seconds + 10 stream = self._open_exec(request, workspace, grpc_deadline) - stdout_parts: list[bytes] = [] stderr_parts: list[bytes] = [] exit_code: int | None = None - - 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) - + try: + 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) + finally: + stream.close() 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"), @@ -1383,7 +1436,7 @@ def _open_exec( request: openshell_pb2.ExecSandboxRequest, workspace: str, timeout: float, - ) -> Iterator[openshell_pb2.ExecSandboxEvent]: + ) -> Generator[openshell_pb2.ExecSandboxEvent, None, None]: """Open ExecSandbox on the owner replica, retrying once unrouted.""" owner = self._owner_of(workspace, request.sandbox) if owner: @@ -1401,9 +1454,20 @@ def _open_exec( else: if first is not None: yield first - yield from events + try: + yield from events + finally: + events.cancel() return - yield from self._stub.ExecSandbox(request, timeout=timeout) + finally: + events.cancel() + events = self._stub.ExecSandbox(request, timeout=timeout) + try: + yield from events + finally: + cancel = getattr(events, "cancel", None) + if cancel is not None: + cancel() def exec( self, @@ -1498,6 +1562,8 @@ def create( gpu_count: int | None = None, gpu: bool = False, driver_config: Mapping[str, Any] | None = None, + timeout: float | None = None, + request_id: str | None = None, ) -> openshell_pb2.SandboxWorkloadTemplate: if template is None: if name is None or not name.strip(): @@ -1525,27 +1591,24 @@ def create( raise SandboxError( "template cannot be combined with template builder fields" ) - response = self._stub.CreateSandboxTemplate( openshell_pb2.CreateSandboxTemplateRequest( workspace_scope=_workspace_scope(workspace), template=template, + request_id=request_id or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return response.template def get( - self, - name: str, - *, - workspace: str, + self, name: str, *, workspace: str, timeout: float | None = None ) -> openshell_pb2.SandboxWorkloadTemplate: response = self._stub.GetSandboxTemplate( openshell_pb2.GetSandboxTemplateRequest( workspace_scope=_workspace_scope(workspace), name=name ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return response.template @@ -1556,6 +1619,7 @@ def list( page_size: int = 100, page_token: str = "", label_selector: str = "", + timeout: float | None = None, ) -> Pager[openshell_pb2.SandboxWorkloadTemplate]: def fetch(token: str) -> Page[openshell_pb2.SandboxWorkloadTemplate]: response = self._stub.ListSandboxTemplates( @@ -1565,7 +1629,7 @@ def fetch(token: str) -> Page[openshell_pb2.SandboxWorkloadTemplate]: page_token=token, label_selector=label_selector, ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return Page( items=list(response.templates), @@ -1595,6 +1659,7 @@ def list_for_all_workspaces( page_size: int = 100, page_token: str = "", label_selector: str = "", + timeout: float | None = None, ) -> Pager[openshell_pb2.SandboxWorkloadTemplate]: def fetch(token: str) -> Page[openshell_pb2.SandboxWorkloadTemplate]: response = self._stub.ListSandboxTemplates( @@ -1604,7 +1669,7 @@ def fetch(token: str) -> Page[openshell_pb2.SandboxWorkloadTemplate]: page_token=token, label_selector=label_selector, ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return Page( items=list(response.templates), @@ -1627,15 +1692,22 @@ def list_all_for_all_workspaces( ).all() def delete( - self, name: str, *, workspace: str, allow_missing: bool = False + self, + name: str, + *, + workspace: str, + allow_missing: bool = False, + timeout: float | None = None, + request_id: str | None = None, ) -> DeletionResult: response = self._stub.DeleteSandboxTemplate( openshell_pb2.DeleteSandboxTemplateRequest( workspace_scope=_workspace_scope(workspace), name=name, allow_missing=allow_missing, + request_id=request_id or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return DeletionResult(DeletionOutcome(response.outcome)) @@ -1672,20 +1744,23 @@ def create( name: str, *, labels: Mapping[str, str] | None = None, + timeout: float | None = None, + request_id: str | None = None, ) -> WorkspaceRef: response = self._stub.CreateWorkspace( openshell_pb2.CreateWorkspaceRequest( name=name, labels=dict(labels) if labels else {}, + request_id=request_id or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return _workspace_ref(response.workspace) - def get(self, name: str) -> WorkspaceRef: + def get(self, name: str, *, timeout: float | None = None) -> WorkspaceRef: response = self._stub.GetWorkspace( openshell_pb2.GetWorkspaceRequest(name=name), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return _workspace_ref(response.workspace) @@ -1695,6 +1770,7 @@ def list( page_size: int = 100, page_token: str = "", label_selector: str | None = None, + timeout: float | None = None, ) -> Pager[WorkspaceRef]: def fetch(token: str) -> Page[WorkspaceRef]: response = self._stub.ListWorkspaces( @@ -1703,7 +1779,7 @@ def fetch(token: str) -> Page[WorkspaceRef]: page_token=token, label_selector=label_selector or "", ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return Page( items=[_workspace_ref(ws) for ws in response.workspaces], @@ -1725,12 +1801,19 @@ def list_all( label_selector=label_selector, ).all() - def delete(self, name: str, *, allow_missing: bool = False) -> DeletionResult: + def delete( + self, + name: str, + *, + allow_missing: bool = False, + timeout: float | None = None, + request_id: str | None = None, + ) -> DeletionResult: response = self._stub.DeleteWorkspace( openshell_pb2.DeleteWorkspaceRequest( - name=name, allow_missing=allow_missing + name=name, allow_missing=allow_missing, request_id=request_id or "" ), - timeout=self._timeout, + timeout=self._timeout if timeout is None else timeout, ) return DeletionResult(DeletionOutcome(response.outcome)) diff --git a/python/wheel_verifier_test.py b/python/wheel_verifier_test.py index 0ade8cd918..a4da3dbd45 100644 --- a/python/wheel_verifier_test.py +++ b/python/wheel_verifier_test.py @@ -29,6 +29,8 @@ def _wheel_files() -> set[str]: files = { "openshell/__init__.py", "openshell/sandbox.py", + "openshell/aio.py", + "openshell/_async_runtime.py", "openshell/py.typed", "openshell/_proto/__init__.py", } diff --git a/tasks/scripts/generate_python_async.py b/tasks/scripts/generate_python_async.py new file mode 100644 index 0000000000..f24d7793c2 --- /dev/null +++ b/tasks/scripts/generate_python_async.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generate native async SDK methods from the synchronous operation bodies. + +This deliberately small source transform knows the SDK's I/O boundaries: +stub calls, pager fetches, stream iteration, and polling sleeps. Transport and +resource ownership are handwritten in _async_runtime.py. No reflection or +dynamic method binding runs in the installed SDK. +""" + +from __future__ import annotations + +import ast +from copy import deepcopy +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +CLASSES = { + "SandboxClient", + "SandboxSession", + "SandboxTemplateClient", + "WorkspaceClient", + "ProviderClient", +} +STREAM_RPCS = {"ExecSandbox", "WatchSandbox"} +STREAM_METHODS = {"exec_stream", "watch"} +FACTORIES = { + "__init__", + "from_active_cluster", + "from_channel", + "from_sandbox_client", + "providers", + "sandbox_templates", + "workspaces", + "list", + "list_for_all_workspaces", + "list_services", + "list_providers", + "list_policies", + "_rpc_timeout", +} + + +def dotted(node: ast.AST) -> str: + return ast.unparse(node) + + +class AsyncMethods(ast.NodeTransformer): + def __init__(self, asynchronous: set[str]) -> None: + self.asynchronous = asynchronous + + def visit_Name(self, node: ast.Name): + if node.id in CLASSES or node.id == "Pager": + node.id = "Async" + node.id + return node + + def visit_Subscript(self, node: ast.Subscript): + if isinstance(node.value, ast.Name) and node.value.id == "Iterator": + return ast.Subscript( + ast.Name("AsyncGenerator", ast.Load()), + ast.Tuple([self.visit(node.slice), ast.Constant(None)], ast.Load()), + ast.Load(), + ) + return self.generic_visit(node) + + def visit_Attribute(self, node: ast.Attribute): + if dotted(node) == "grpc.Channel": + return ast.parse("grpc.aio.Channel", mode="eval").body + return self.generic_visit(node) + + def visit_ImportFrom(self, node: ast.ImportFrom): + # ProviderClient is defined in this generated module. + if node.module == "providers": + return None + return node + + def visit_Call(self, node: ast.Call): + original = dotted(node.func) + self.generic_visit(node) + if original == "_error_mapping_channel": + return node.args[0] + if original.startswith("self._stub."): + if original.rsplit(".", 1)[1] in STREAM_RPCS: + return node + return ast.Await(ast.Call(ast.Name("_await_rpc", ast.Load()), [node], [])) + if original == "time.sleep": + node.func = ast.parse("asyncio.sleep", mode="eval").body + return ast.Await(node) + if original == "stream.close": + node.func = ast.parse("stream.aclose", mode="eval").body + return ast.Await(node) + if isinstance(node.func, ast.Attribute) and ( + node.func.attr == "all" + or ( + node.func.attr in self.asynchronous + and original.startswith(("self.", "client.")) + ) + ): + return ast.Await(node) + return node + + def visit_For(self, node: ast.For): + original = dotted(node.iter) + self.generic_visit(node) + if original == "stream" or any( + f".{method}(" in original for method in STREAM_METHODS + ): + return ast.AsyncFor( + node.target, node.iter, node.body, node.orelse, node.type_comment + ) + return node + + def visit_FunctionDef(self, node: ast.FunctionDef): + self.generic_visit(node) + if ( + node.name in self.asynchronous + or node.name in STREAM_METHODS + or node.name == "fetch" + ): + # Stream failures occur during iteration; catch them at the method + # boundary without swallowing task cancellation. + if node.name in STREAM_METHODS: + doc = ( + node.body[:1] + if isinstance(node.body[0], ast.Expr) + and isinstance(node.body[0].value, ast.Constant) + else [] + ) + handler = ast.parse( + "try:\n pass\nexcept grpc.RpcError as error:\n raise from_grpc_error(error) from error" + ).body[0] + assert isinstance(handler, ast.Try) + handler.body = node.body[len(doc) :] + node.body = [*doc, handler] + return ast.AsyncFunctionDef( + name=node.name, + args=node.args, + body=node.body, + decorator_list=node.decorator_list, + returns=node.returns, + type_comment=node.type_comment, + ) + return node + + +def generate(root: Path = ROOT) -> str: + sources = [ + ast.parse((root / f"python/openshell/{name}.py").read_text()) + for name in ("sandbox", "providers") + ] + classes = [ + deepcopy(node) + for source in sources + for node in source.body + if isinstance(node, ast.ClassDef) and node.name in CLASSES + ] + for cls in classes: + if cls.name == "SandboxClient": + # Native call metadata and owner routing differ between grpc and + # grpc.aio; the async runtime owns these two implementations. + cls.body = [ + node + for node in cls.body + if not isinstance(node, ast.FunctionDef) + or node.name not in {"_owner_of", "_open_exec"} + ] + asynchronous = { + method.name + for cls in classes + for method in cls.body + if isinstance(method, ast.FunctionDef) + and method.name not in FACTORIES | STREAM_METHODS + and any( + isinstance(node, ast.Call) and dotted(node.func).startswith("self._stub.") + for node in ast.walk(method) + ) + } + asynchronous |= { + "exec", + "exec_python", + "get_session", + "create_session", + "create_session_from_template", + "wait_deleted", + "_wait_for_phase", + "wait_completed", + } + # Propagate async through helper calls and pager collection. + while True: + previous = asynchronous.copy() + for cls in classes: + for method in cls.body: + if ( + not isinstance(method, ast.FunctionDef) + or method.name in FACTORIES | STREAM_METHODS + ): + continue + if any( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and (node.func.attr in asynchronous or node.func.attr == "all") + for node in ast.walk(method) + ): + asynchronous.add(method.name) + if previous == asynchronous: + break + result = [] + for cls in classes: + if cls.name == "SandboxClient": + cls.bases = [ast.Name("AsyncTransport", ast.Load())] + cls.body = [ + node + for node in cls.body + if not isinstance(node, ast.FunctionDef) + or node.name + not in {"__init__", "close", "__enter__", "__exit__", "from_channel"} + ] + cls.name = "Async" + cls.name + result.append(AsyncMethods(asynchronous).visit(cls)) + tree = ast.fix_missing_locations(ast.Module(result, type_ignores=[])) + used = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)} + imports = [ + "from __future__ import annotations", + "import asyncio", + "import grpc.aio", + "from collections.abc import AsyncGenerator", + "from ._async_runtime import AsyncPager, AsyncTransport, _await_rpc", + "from .errors import from_grpc_error", + ] + for source, module in zip(sources, ("sandbox", "providers"), strict=True): + for node in source.body: + if isinstance(node, (ast.Import, ast.ImportFrom)) and not ( + isinstance(node, ast.ImportFrom) and node.module == "__future__" + ): + imports.append(ast.unparse(node)) + if isinstance(node, ast.If) and dotted(node.test) == "TYPE_CHECKING": + for child in node.body: + if isinstance(child, (ast.Import, ast.ImportFrom)) and not ( + isinstance(child, ast.ImportFrom) + and child.module in {"providers", "sandbox"} + ): + imports.append(ast.unparse(child)) + definitions = { + node.name + for node in source.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) + } | { + target.id + for node in source.body + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Name) + } + names = sorted((definitions & used) - CLASSES - {"Pager", "T"}) + if names: + imports.append(f"from .{module} import ({', '.join(names)})") + header = "# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n# SPDX-License-Identifier: Apache-2.0\n\n# Generated by tasks/scripts/generate_python_async.py. Do not edit.\n" + return ( + header + "\n".join(dict.fromkeys(imports)) + "\n\n" + ast.unparse(tree) + "\n" + ) + + +if __name__ == "__main__": + (ROOT / "python/openshell/aio.py").write_text(generate()) diff --git a/tasks/scripts/generate_python_proto.py b/tasks/scripts/generate_python_proto.py index e0a1896d0d..ba0b49b15e 100644 --- a/tasks/scripts/generate_python_proto.py +++ b/tasks/scripts/generate_python_proto.py @@ -9,6 +9,8 @@ import sys from pathlib import Path +from generate_python_async import generate + PROTO_FILES = [ "proto/openshell.proto", "proto/datamodel.proto", @@ -67,6 +69,10 @@ def main() -> None: text = re.sub(pattern, r"from . import \1\2", text, flags=re.MULTILINE) file_path.write_text(text) + # Use the checkout's authored methods, but write into the requested build + # tree just like protoc. The generator has no runtime SDK imports. + (output_dir.parent / "aio.py").write_text(generate()) + if __name__ == "__main__": main() diff --git a/tasks/scripts/verify-python-wheel.py b/tasks/scripts/verify-python-wheel.py index b145d1fd48..5d84cc3d78 100644 --- a/tasks/scripts/verify-python-wheel.py +++ b/tasks/scripts/verify-python-wheel.py @@ -39,6 +39,8 @@ def _required_files() -> set[str]: files = { "openshell/__init__.py", "openshell/sandbox.py", + "openshell/aio.py", + "openshell/_async_runtime.py", "openshell/py.typed", "openshell/_proto/__init__.py", } From 5b093f792f3f55b7367e200873dc5744fe1ccfb2 Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Mon, 2 Nov 2026 16:31:22 -0500 Subject: [PATCH 2/2] feat(python): expose sandbox services and structured logs Signed-off-by: Max Dubrinsky --- docs/sdk/python.mdx | 30 +++ python/openshell/async_client_test.py | 106 +++++++++ python/openshell/sandbox.py | 297 +++++++++++++++++++++++-- python/openshell/types.py | 56 +++++ tasks/scripts/generate_python_async.py | 21 +- 5 files changed, 495 insertions(+), 15 deletions(-) create mode 100644 python/openshell/types.py diff --git a/docs/sdk/python.mdx b/docs/sdk/python.mdx index 9cea0f5085..c3968e9446 100644 --- a/docs/sdk/python.mdx +++ b/docs/sdk/python.mdx @@ -187,6 +187,36 @@ Async task cancellation propagates to the active RPC. For early exit from `exec_stream()`, use `contextlib.aclosing()` so the stream is cancelled promptly. Gateway RPC failures raise `GatewayError`, preserving the gRPC status and details. +## Manage Services and Read Logs + +Both sandbox clients expose service creation, lookup, pagination, and deletion. +Use `get_logs()` for structured gateway and supervisor logs: + +```python +from openshell import SandboxClient + +with SandboxClient.from_active_cluster() as client: + session = client.get_session("worker", workspace="default") + endpoint = session.expose_service(name="api", target_port=8000) + print(endpoint.url) + for page in session.list_services(): + for service in page.items: + print(service.endpoint.name, service.url) + logs = session.get_logs(lines=100, sources=["sandbox"]) + for entry in logs.logs: + print(entry.message, dict(entry.fields)) + session.delete_service(name="api", allow_missing=True) +``` + +`get_logs()` preserves `buffer_total`, the unfiltered count of buffered entries. +Its optional `since` argument takes a timezone-aware `datetime`. These logs do +not include arbitrary workload output written to files; use exec to read those +files. Service and log results retain their complete generated protobuf fields. + +Import supported protobuf-backed inputs and results from `openshell.types`, +including `SandboxSpec`, `SandboxTemplate`, `SandboxPolicy`, `SandboxPhase`, +`ServiceEndpointResponse`, and `GetSandboxLogsResponse`. + ## Call Generated RPCs The generated gRPC stubs include concrete request and response types for synchronous `grpc.Channel` and asynchronous `grpc.aio.Channel` connections. Install the dependency type stubs in your development environment for editor signatures and type checking: diff --git a/python/openshell/async_client_test.py b/python/openshell/async_client_test.py index f501f16e06..941d608731 100644 --- a/python/openshell/async_client_test.py +++ b/python/openshell/async_client_test.py @@ -82,6 +82,51 @@ async def ExecSandbox(self, request, context): event.exit.exit_code = 0 yield event + async def ExposeService(self, request, context): + self.requests.append((request, context.invocation_metadata())) + return pb.ServiceEndpointResponse( + endpoint=pb.ServiceEndpoint( + name=request.name, + sandbox=request.sandbox, + target_port=request.target_port, + authorization_mode=request.authorization_mode, + ), + url="https://service.example/", + ) + + async def GetService(self, request, context): + self.requests.append((request, context.invocation_metadata())) + return pb.ServiceEndpointResponse( + endpoint=pb.ServiceEndpoint(name=request.name), + url="https://service.example/", + ) + + async def ListServices(self, request, context): + self.requests.append((request, context.invocation_metadata())) + return pb.ListServicesResponse( + services=[ + pb.ServiceEndpointResponse( + endpoint=pb.ServiceEndpoint(name=request.page_token or "first") + ) + ], + next_page_token="second" if not request.page_token else "", + ) + + async def DeleteService(self, request, context): + self.requests.append((request, context.invocation_metadata())) + return pb.DeleteServiceResponse(outcome=pb.DELETION_OUTCOME_ALREADY_ABSENT) + + async def GetSandboxLogs(self, request, context): + self.requests.append((request, context.invocation_metadata())) + return pb.GetSandboxLogsResponse( + logs=[ + pb.SandboxLogLine( + message="denied", fields={"action": "deny"}, source="sandbox" + ) + ], + buffer_total=42, + ) + @asynccontextmanager async def gateway(): @@ -203,3 +248,64 @@ def test_async_methods_are_real_coroutines(): assert not inspect.iscoroutinefunction(AsyncSandboxClient.list) client = create_autospec(AsyncSandboxClient, instance=True) assert isinstance(client.create, AsyncMock) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_services_logs_and_session_scope(asynchronous): + from datetime import UTC, datetime + + from openshell import DeletionOutcome, SandboxClient, ServiceAuthorizationMode + + async with gateway() as (service, target): + client = AsyncSandboxClient(target) if asynchronous else SandboxClient(target) + + async def invoke(method, *args, **kwargs): + if asynchronous: + return await method(*args, **kwargs) + return await asyncio.to_thread(method, *args, **kwargs) + + try: + session = await invoke(client.get_session, "worker", workspace="team") + endpoint = await invoke( + session.expose_service, + name="api", + target_port=8080, + authorization_mode=ServiceAuthorizationMode.BEARER_PASSTHROUGH, + request_id="expose-1", + ) + assert endpoint.url == "https://service.example/" + request = service.requests[-1][0] + assert request.workspace_scope.workspace == "team" + assert request.sandbox == "worker" + assert request.request_id == "expose-1" + assert ( + endpoint.endpoint.authorization_mode + == pb.SERVICE_AUTHORIZATION_MODE_BEARER_PASSTHROUGH + ) + assert ( + await invoke(session.get_service, name="api") + ).endpoint.name == "api" + endpoints = await invoke(session.list_all_services) + assert [item.endpoint.name for item in endpoints] == ["first", "second"] + result = await invoke( + session.delete_service, + name="api", + allow_missing=True, + request_id="delete-1", + ) + assert result.outcome == DeletionOutcome.ALREADY_ABSENT + since = datetime(2026, 1, 1, tzinfo=UTC) + logs = await invoke( + session.get_logs, + lines=10, + since=since, + sources=["sandbox"], + min_level="WARN", + ) + assert logs.buffer_total == 42 + assert logs.logs[0].fields["action"] == "deny" + assert service.requests[-1][0].since_time.ToDatetime(tzinfo=UTC) == since + with pytest.raises(ValueError, match="timezone-aware"): + await invoke(session.get_logs, since=datetime(2026, 1, 1)) + finally: + await invoke(client.close) diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index 092ea2506c..c3f283fece 100644 --- a/python/openshell/sandbox.py +++ b/python/openshell/sandbox.py @@ -187,7 +187,8 @@ class _ClientCallDetails(_ClientCallDetailsBase, grpc.ClientCallDetails): if TYPE_CHECKING: import builtins - from collections.abc import Callable, Generator, Iterator, Mapping, Sequence + from collections.abc import Callable, Generator, Mapping, Sequence + from datetime import datetime from .providers import ProviderClient @@ -707,6 +708,129 @@ def resume(self, *, restart_if_restore_fails: bool = False) -> SandboxRef: ) return self.sandbox + def exec_stream( + self, + command: Sequence[str], + *, + workdir: str | None = None, + env: Mapping[str, str] | None = None, + stdin: bytes | None = None, + timeout_seconds: int | None = None, + no_login_shell: bool = False, + timeout: float | None = None, + request_id: str | None = None, + ) -> Generator[ExecChunk | ExecResult, None, None]: + stream = self._client.exec_stream( + self.sandbox.name, + command, + workspace=self._workspace, + workdir=workdir, + env=env, + stdin=stdin, + timeout_seconds=timeout_seconds, + no_login_shell=no_login_shell, + timeout=timeout, + request_id=request_id, + ) + try: + yield from stream + finally: + stream.close() + + def expose_service( + self, + *, + name: str = "", + target_port: int, + domain: bool = False, + authorization_mode: ServiceAuthorizationMode = ServiceAuthorizationMode.STRIP, + request_id: str | None = None, + timeout: float | None = None, + ) -> openshell_pb2.ServiceEndpointResponse: + return self._client.expose_service( + self.sandbox.name, + workspace=self._workspace, + name=name, + target_port=target_port, + domain=domain, + authorization_mode=authorization_mode, + request_id=request_id, + timeout=timeout, + ) + + def get_service( + self, *, name: str = "", timeout: float | None = None + ) -> openshell_pb2.ServiceEndpointResponse: + return self._client.get_service( + self.sandbox.name, workspace=self._workspace, name=name, timeout=timeout + ) + + def list_services( + self, + *, + page_size: int = 100, + page_token: str = "", + timeout: float | None = None, + ) -> Pager[openshell_pb2.ServiceEndpointResponse]: + return self._client.list_services( + self.sandbox.name, + workspace=self._workspace, + page_size=page_size, + page_token=page_token, + timeout=timeout, + ) + + def list_all_services( + self, + *, + page_size: int = 100, + page_token: str = "", + timeout: float | None = None, + ) -> builtins.list[openshell_pb2.ServiceEndpointResponse]: + return self._client.list_all_services( + self.sandbox.name, + workspace=self._workspace, + page_size=page_size, + page_token=page_token, + timeout=timeout, + ) + + def delete_service( + self, + *, + name: str = "", + allow_missing: bool = False, + request_id: str | None = None, + timeout: float | None = None, + ) -> DeletionResult: + return self._client.delete_service( + self.sandbox.name, + workspace=self._workspace, + name=name, + allow_missing=allow_missing, + request_id=request_id, + timeout=timeout, + ) + + def get_logs( + self, + *, + lines: int = 2000, + since: datetime | None = None, + sources: Sequence[str] = (), + min_level: str = "", + timeout: float | None = None, + ) -> openshell_pb2.GetSandboxLogsResponse: + return self._client.get_logs( + self.sandbox.name, + workspace=self._workspace, + lines=lines, + since=since, + sources=sources, + min_level=min_level, + timeout=timeout, + ) + class SandboxClient: """gRPC client for sandbox CRUD and command execution.""" @@ -1368,7 +1492,7 @@ def exec_stream( no_login_shell: bool = False, timeout: float | None = None, request_id: str | None = None, - ) -> Iterator[ExecChunk | ExecResult]: + ) -> Generator[ExecChunk | ExecResult, None, None]: if not command: raise SandboxError("command must not be empty") request = openshell_pb2.ExecSandboxRequest( @@ -1483,7 +1607,7 @@ def exec( no_login_shell: bool = False, ) -> ExecResult: result: ExecResult | None = None - for item in self.exec_stream( + stream = self.exec_stream( sandbox, command, workspace=workspace, @@ -1492,16 +1616,20 @@ def exec( stdin=stdin, timeout_seconds=timeout_seconds, no_login_shell=no_login_shell, - ): - 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 + ) + try: + for item in stream: + 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 + finally: + stream.close() if result is None: raise SandboxError("ExecSandbox did not return a result") return result @@ -1535,6 +1663,144 @@ def exec_python( timeout_seconds=timeout_seconds, ) + def expose_service( + self, + sandbox: str, + *, + workspace: str, + name: str = "", + target_port: int, + domain: bool = False, + authorization_mode: ServiceAuthorizationMode = ServiceAuthorizationMode.STRIP, + request_id: str | None = None, + timeout: float | None = None, + ) -> openshell_pb2.ServiceEndpointResponse: + """Create or update an HTTP endpoint, retaining its routing metadata and URL.""" + return self._stub.ExposeService( + openshell_pb2.ExposeServiceRequest( + workspace_scope=_workspace_scope(workspace), + sandbox=sandbox, + name=name, + target_port=target_port, + domain=domain, + authorization_mode=f"SERVICE_AUTHORIZATION_MODE_{authorization_mode.name}", + request_id=request_id or "", + ), + timeout=self._timeout if timeout is None else timeout, + ) + + def get_service( + self, + sandbox: str, + *, + workspace: str, + name: str = "", + timeout: float | None = None, + ) -> openshell_pb2.ServiceEndpointResponse: + """Fetch an endpoint. An empty name selects the unnamed service.""" + return self._stub.GetService( + openshell_pb2.GetServiceRequest( + workspace_scope=_workspace_scope(workspace), + sandbox=sandbox, + name=name, + ), + timeout=self._timeout if timeout is None else timeout, + ) + + def list_services( + self, + sandbox: str = "", + *, + workspace: str, + page_size: int = 100, + page_token: str = "", + timeout: float | None = None, + ) -> Pager[openshell_pb2.ServiceEndpointResponse]: + """List one sandbox's endpoints, or every endpoint in the named workspace.""" + scope = _workspace_scope(workspace) + + def fetch(token: str) -> Page[openshell_pb2.ServiceEndpointResponse]: + response = self._stub.ListServices( + openshell_pb2.ListServicesRequest( + workspace_scope=scope, + sandbox=sandbox, + page_size=page_size, + page_token=token, + ), + timeout=self._timeout if timeout is None else timeout, + ) + return Page(list(response.services), response.next_page_token) + + return Pager(fetch, page_token) + + def list_all_services( + self, + sandbox: str = "", + *, + workspace: str, + page_size: int = 100, + page_token: str = "", + timeout: float | None = None, + ) -> builtins.list[openshell_pb2.ServiceEndpointResponse]: + return self.list_services( + sandbox, + workspace=workspace, + page_size=page_size, + page_token=page_token, + timeout=timeout, + ).all() + + def delete_service( + self, + sandbox: str, + *, + workspace: str, + name: str = "", + allow_missing: bool = False, + request_id: str | None = None, + timeout: float | None = None, + ) -> DeletionResult: + response = self._stub.DeleteService( + openshell_pb2.DeleteServiceRequest( + workspace_scope=_workspace_scope(workspace), + sandbox=sandbox, + name=name, + allow_missing=allow_missing, + request_id=request_id or "", + ), + timeout=self._timeout if timeout is None else timeout, + ) + return DeletionResult(DeletionOutcome(response.outcome)) + + def get_logs( + self, + sandbox: str, + *, + workspace: str, + lines: int = 2000, + since: datetime | None = None, + sources: Sequence[str] = (), + min_level: str = "", + timeout: float | None = None, + ) -> openshell_pb2.GetSandboxLogsResponse: + """Fetch correlated gateway/supervisor logs and the unfiltered buffer size. + + These are not arbitrary workload stdout. ``since`` must be timezone-aware. + """ + request = openshell_pb2.GetSandboxLogsRequest( + workspace_scope=_workspace_scope(workspace), + sandbox=sandbox, + lines=lines, + sources=sources, + min_level=min_level, + ) + if since is not None: + _require_aware_datetime(since) + request.since_time.FromDatetime(since) + return self._stub.GetSandboxLogs( + request, timeout=self._timeout if timeout is None else timeout + ) + class SandboxTemplateClient: """gRPC client for reusable sandbox template lifecycle operations.""" @@ -2745,3 +3011,8 @@ def _resolve_active_cluster() -> str: if value == "": raise SandboxError("no active gateway configured") return value + + +def _require_aware_datetime(value: datetime) -> None: + if value.tzinfo is None or value.utcoffset() is None: + raise ValueError("timestamp must be timezone-aware") diff --git a/python/openshell/types.py b/python/openshell/types.py new file mode 100644 index 0000000000..e2e0a28922 --- /dev/null +++ b/python/openshell/types.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public imports for the protobuf-backed SDK inputs and complete RPC results. + +These are the generated message classes, not independent copies of the schema. +Protobuf presence, enum, and oneof APIs remain available. Use bindings from the +same release as the gateway. Do not log messages containing credential material. +""" + +from ._proto.datamodel_pb2 import ObjectMeta, Provider +from ._proto.openshell_pb2 import ( + AttachSandboxProviderResponse, + DetachSandboxProviderResponse, + GetSandboxLogsResponse, + GetSandboxPolicyStatusResponse, + ProviderMutationReceipt, + ProviderReadinessState, + ProviderReadinessStatus, + ResourceRequirements, + SandboxLogLine, + SandboxPhase, + SandboxPolicyRevision, + SandboxSpec, + SandboxStreamEvent, + SandboxTemplate, + SandboxWorkloadTemplate, + ServiceEndpointResponse, + UpdateConfigResponse, +) +from ._proto.sandbox_pb2 import GetSandboxConfigResponse, SandboxPolicy, SettingValue + +__all__ = [ + "AttachSandboxProviderResponse", + "DetachSandboxProviderResponse", + "GetSandboxConfigResponse", + "GetSandboxLogsResponse", + "GetSandboxPolicyStatusResponse", + "ObjectMeta", + "Provider", + "ProviderMutationReceipt", + "ProviderReadinessState", + "ProviderReadinessStatus", + "ResourceRequirements", + "SandboxLogLine", + "SandboxPhase", + "SandboxPolicy", + "SandboxPolicyRevision", + "SandboxSpec", + "SandboxStreamEvent", + "SandboxTemplate", + "SandboxWorkloadTemplate", + "ServiceEndpointResponse", + "SettingValue", + "UpdateConfigResponse", +] diff --git a/tasks/scripts/generate_python_async.py b/tasks/scripts/generate_python_async.py index f24d7793c2..a855bc2689 100644 --- a/tasks/scripts/generate_python_async.py +++ b/tasks/scripts/generate_python_async.py @@ -56,10 +56,16 @@ def visit_Name(self, node: ast.Name): return node def visit_Subscript(self, node: ast.Subscript): - if isinstance(node.value, ast.Name) and node.value.id == "Iterator": + if isinstance(node.value, ast.Name) and node.value.id in { + "Iterator", + "Generator", + }: + item_type = ( + node.slice.elts[0] if isinstance(node.slice, ast.Tuple) else node.slice + ) return ast.Subscript( ast.Name("AsyncGenerator", ast.Load()), - ast.Tuple([self.visit(node.slice), ast.Constant(None)], ast.Load()), + ast.Tuple([self.visit(item_type), ast.Constant(None)], ast.Load()), ast.Load(), ) return self.generic_visit(node) @@ -111,6 +117,17 @@ def visit_For(self, node: ast.For): ) return node + def visit_Expr(self, node: ast.Expr): + if isinstance(node.value, ast.YieldFrom): + return ast.AsyncFor( + ast.Name("item", ast.Store()), + self.visit(node.value.value), + [ast.Expr(ast.Yield(ast.Name("item", ast.Load())))], + [], + None, + ) + return self.generic_visit(node) + def visit_FunctionDef(self, node: ast.FunctionDef): self.generic_visit(node) if (