From 9168115c96f4c3ce4d4322c8a52041574f2e9ca0 Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Mon, 2 Nov 2026 16:33:16 -0500 Subject: [PATCH] feat(python): add sandbox watches and completion waits Signed-off-by: Max Dubrinsky --- docs/sdk/python.mdx | 34 +++++++ python/openshell/async_client_test.py | 91 ++++++++++++++++- python/openshell/sandbox.py | 141 ++++++++++++++++++++++++++ 3 files changed, 264 insertions(+), 2 deletions(-) diff --git a/docs/sdk/python.mdx b/docs/sdk/python.mdx index c3968e9446..4e74fe0728 100644 --- a/docs/sdk/python.mdx +++ b/docs/sdk/python.mdx @@ -217,6 +217,40 @@ Import supported protobuf-backed inputs and results from `openshell.types`, including `SandboxSpec`, `SandboxTemplate`, `SandboxPolicy`, `SandboxPhase`, `ServiceEndpointResponse`, and `GetSandboxLogsResponse`. +## Observe a Sandbox + +`watch()` streams status snapshots and optionally logs and platform events. It +returns complete `SandboxStreamEvent` messages, including warnings and draft +policy notifications. Use `contextlib.closing()` for sync iteration or +`contextlib.aclosing()` for async iteration to cancel promptly on early exit: + +```python +from contextlib import aclosing + +async def observe(client): + async with aclosing(client.watch("worker", workspace="default", follow_logs=True)) as events: + async for event in events: + kind = event.WhichOneof("payload") + if kind == "log": + print(event.log.message) + elif kind == "warning": + print(event.warning.message) +``` + +Watch streams have no deadline unless `timeout` is supplied. To resume logs and +platform events, retain the greatest nonempty cursor you have processed on that +stream and pass it as `resume_after_cursor`. Status and warning messages may +have no cursor. The SDK does not reconnect or reset cursors automatically; +`OUT_OF_RANGE` means the cursor is no longer usable and requires an explicit +decision to start a new watch. Preserve warnings as evidence of event gaps. + +Use `wait_completed()` for a job running as the sandbox's canonical main process. +It returns the completed `SandboxRef`, raises on error, stopping, or deletion, +and detects a same-name replacement. Supply `expected_sandbox_id=sandbox.id` +from creation to pin the target before the first poll. Its `timeout_seconds` +bounds polling RPCs and sleeps; `poll_interval` controls the polling frequency. +Use `exec()` to wait for a separately executed command instead. + ## 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 941d608731..03a9eabcad 100644 --- a/python/openshell/async_client_test.py +++ b/python/openshell/async_client_test.py @@ -28,6 +28,7 @@ class Gateway(rpc.OpenShellServicer): def __init__(self): self.requests = [] + self.snapshots = [] self.started = asyncio.Event() self.cancelled = asyncio.Event() @@ -41,14 +42,19 @@ async def GetSandbox(self, request, context): await asyncio.Future() finally: self.cancelled.set() + identity, phase = ( + self.snapshots.pop(0) + if self.snapshots + else ("identity", pb.SANDBOX_PHASE_READY) + ) return pb.SandboxResponse( sandbox=pb.Sandbox( metadata=dm.ObjectMeta( - id="identity", + id=identity, name=request.name, workspace=request.workspace_scope.workspace, ), - status=pb.SandboxStatus(phase=pb.SANDBOX_PHASE_READY, exit_code=0), + status=pb.SandboxStatus(phase=phase, exit_code=0), ) ) @@ -127,6 +133,23 @@ async def GetSandboxLogs(self, request, context): buffer_total=42, ) + async def WatchSandbox(self, request, context): + self.requests.append((request, context.invocation_metadata())) + if request.resume_after_cursor == "expired": + await context.abort(grpc.StatusCode.OUT_OF_RANGE, "cursor expired") + yield pb.SandboxStreamEvent( + log=pb.SandboxLogLine(message="log"), cursor="cursor-1" + ) + yield pb.SandboxStreamEvent( + warning=pb.SandboxStreamWarning(message="events dropped") + ) + if request.sandbox == "blocked": + self.started.set() + try: + await asyncio.Future() + finally: + self.cancelled.set() + @asynccontextmanager async def gateway(): @@ -309,3 +332,67 @@ async def invoke(method, *args, **kwargs): await invoke(session.get_logs, since=datetime(2026, 1, 1)) finally: await invoke(client.close) + + +async def test_watch_preserves_cursors_warnings_and_cancels(): + async with ( + gateway() as (service, target), + AsyncSandboxClient(target, bearer_token="watch-token") as client, + ): + stream = client.watch( + "blocked", + workspace="team", + follow_logs=True, + resume_after_cursor="previous", + ) + event = await anext(stream) + assert event.cursor == "cursor-1" + warning = await anext(stream) + assert warning.WhichOneof("payload") == "warning" + assert warning.cursor == "" + await asyncio.wait_for(service.started.wait(), 2) + await stream.aclose() + await asyncio.wait_for(service.cancelled.wait(), 2) + request, metadata = service.requests[-1] + assert request.resume_after_cursor == "previous" + assert ("authorization", "Bearer watch-token") in metadata + with pytest.raises(GatewayError) as failure: + async for _ in client.watch( + "worker", workspace="team", resume_after_cursor="expired" + ): + pass + assert failure.value.code() == grpc.StatusCode.OUT_OF_RANGE + assert len(service.requests) == 2 # No hidden reconnect or cursor reset. + + +async def test_wait_completed_pins_identity_and_checks_terminal_outcomes(): + async with gateway() as (service, target), AsyncSandboxClient(target) as client: + service.snapshots = [ + ("one", pb.SANDBOX_PHASE_READY), + ("one", pb.SANDBOX_PHASE_COMPLETED), + ] + completed = await client.wait_completed( + "worker", workspace="team", timeout_seconds=1, poll_interval=0.001 + ) + assert completed.id == "one" + assert completed.status.exit_code == 0 + service.snapshots = [("replacement", pb.SANDBOX_PHASE_COMPLETED)] + with pytest.raises(SandboxError, match="replaced"): + await client.wait_completed( + "worker", workspace="team", expected_sandbox_id="one" + ) + for phase in ( + pb.SANDBOX_PHASE_ERROR, + pb.SANDBOX_PHASE_STOPPED, + pb.SANDBOX_PHASE_DELETING, + ): + service.snapshots = [("one", phase)] + with pytest.raises(SandboxError, match="before completion"): + await client.wait_completed("worker", workspace="team") + with pytest.raises(ValueError, match="poll_interval"): + await client.wait_completed("worker", workspace="team", poll_interval=0) + with pytest.raises(GatewayError) as failure: + await client.wait_completed( + "blocked", workspace="team", timeout_seconds=0.05 + ) + assert failure.value.code() == grpc.StatusCode.DEADLINE_EXCEEDED diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index c3f283fece..c40b4712fb 100644 --- a/python/openshell/sandbox.py +++ b/python/openshell/sandbox.py @@ -831,6 +831,56 @@ def get_logs( timeout=timeout, ) + def watch( + self, + *, + follow_status: bool = True, + follow_logs: bool = False, + follow_events: bool = False, + log_tail_lines: int = 0, + event_tail: int = 0, + stop_on_terminal: bool = False, + since: datetime | None = None, + log_sources: Sequence[str] = (), + log_min_level: str = "", + resume_after_cursor: str = "", + timeout: float | None = None, + ) -> Generator[openshell_pb2.SandboxStreamEvent, None, None]: + stream = self._client.watch( + self.sandbox.name, + workspace=self._workspace, + follow_status=follow_status, + follow_logs=follow_logs, + follow_events=follow_events, + log_tail_lines=log_tail_lines, + event_tail=event_tail, + stop_on_terminal=stop_on_terminal, + since=since, + log_sources=log_sources, + log_min_level=log_min_level, + resume_after_cursor=resume_after_cursor, + timeout=timeout, + ) + try: + yield from stream + finally: + stream.close() + + def wait_completed( + self, + *, + timeout_seconds: float = 300.0, + poll_interval: float = 1.0, + expected_sandbox_id: str | None = None, + ) -> SandboxRef: + return self._client.wait_completed( + self.sandbox.name, + workspace=self._workspace, + timeout_seconds=timeout_seconds, + poll_interval=poll_interval, + expected_sandbox_id=expected_sandbox_id, + ) + class SandboxClient: """gRPC client for sandbox CRUD and command execution.""" @@ -1801,6 +1851,97 @@ def get_logs( request, timeout=self._timeout if timeout is None else timeout ) + def watch( + self, + sandbox: str, + *, + workspace: str, + follow_status: bool = True, + follow_logs: bool = False, + follow_events: bool = False, + log_tail_lines: int = 0, + event_tail: int = 0, + stop_on_terminal: bool = False, + since: datetime | None = None, + log_sources: Sequence[str] = (), + log_min_level: str = "", + resume_after_cursor: str = "", + timeout: float | None = None, + ) -> Generator[openshell_pb2.SandboxStreamEvent, None, None]: + """Stream complete status, log, event, warning, and draft-policy messages. + + No deadline is imposed unless ``timeout`` is supplied. The caller owns + reconnects and cursor persistence; OUT_OF_RANGE is never silently reset. + Close this iterator on early exit to cancel the RPC. + """ + request = openshell_pb2.WatchSandboxRequest( + workspace_scope=_workspace_scope(workspace), + sandbox=sandbox, + follow_status=follow_status, + follow_logs=follow_logs, + follow_events=follow_events, + log_tail_lines=log_tail_lines, + event_tail=event_tail, + stop_on_terminal=stop_on_terminal, + log_sources=log_sources, + log_min_level=log_min_level, + resume_after_cursor=resume_after_cursor, + ) + if since is not None: + _require_aware_datetime(since) + request.since_time.FromDatetime(since) + stream = self._stub.WatchSandbox(request, timeout=timeout) + try: + yield from stream + finally: + stream.cancel() + + def wait_completed( + self, + name: str, + *, + workspace: str, + timeout_seconds: float = 300.0, + poll_interval: float = 1.0, + expected_sandbox_id: str | None = None, + ) -> SandboxRef: + """Wait for the canonical main process to complete successfully. + + ERROR, STOPPED, DELETING, or a same-name replacement raise SandboxError. + Suspension does not imply completion. Exec child processes are not + observed by this method. Every RPC and sleep is bounded by the remaining + overall timeout. Pass the create result's ID to pin identity immediately. + """ + if not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + raise ValueError("timeout_seconds must be finite and greater than zero") + if not math.isfinite(poll_interval) or poll_interval <= 0: + raise ValueError("poll_interval must be finite and greater than zero") + deadline = time.monotonic() + timeout_seconds + identity = expected_sandbox_id + while time.monotonic() < deadline: + sandbox = self.get( + name, + workspace=workspace, + timeout=max(0.0, min(self._timeout, deadline - time.monotonic())), + ) + if identity is not None and sandbox.id != identity: + raise SandboxError( + f"sandbox {name} was replaced while waiting for completion" + ) + identity = sandbox.id + if sandbox.phase == openshell_pb2.SANDBOX_PHASE_COMPLETED: + return sandbox + if sandbox.phase in ( + openshell_pb2.SANDBOX_PHASE_ERROR, + openshell_pb2.SANDBOX_PHASE_STOPPED, + openshell_pb2.SANDBOX_PHASE_DELETING, + ): + raise SandboxError( + f"sandbox {name} entered {openshell_pb2.SandboxPhase.Name(sandbox.phase)} before completion" + ) + time.sleep(max(0.0, min(poll_interval, deadline - time.monotonic()))) + raise SandboxError(f"sandbox {name} did not complete within timeout") + class SandboxTemplateClient: """gRPC client for reusable sandbox template lifecycle operations."""