From e57548c7ed311e043341c017285a12117acab26d Mon Sep 17 00:00:00 2001 From: divesh Date: Sun, 25 Oct 2026 22:09:43 -0700 Subject: [PATCH 1/2] feat(sdk): route Python SDK exec to the sandbox owner replica Ask GetSandbox for the owner replica, cache it for 30 seconds, and send it as x-openshell-replica on ExecSandbox. Retry once without the hint when the per-replica route has no healthy endpoint, forgetting the cached owner. Part of #4223 Signed-off-by: divesh --- python/openshell/sandbox.py | 99 +++++++++++++++++++++++++++- python/openshell/sandbox_test.py | 109 ++++++++++++++++++++++++++++++- 2 files changed, 204 insertions(+), 4 deletions(-) diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index 2eaf6a64fa..d9ac242f23 100644 --- a/python/openshell/sandbox.py +++ b/python/openshell/sandbox.py @@ -29,7 +29,7 @@ openshell_pb2, openshell_pb2_grpc, ) -from .errors import _error_mapping_channel +from .errors import _error_mapping_channel, from_grpc_error from .mutations import DeletionOutcome, DeletionResult _ClientCallDetailsBase = namedtuple( @@ -40,6 +40,9 @@ _OAUTH_MAX_RESPONSE_BYTES = 1 << 20 _PAGER_MAX_CONSUMED_TOKENS = 10_000 _PAGER_MAX_CONSUMED_TOKEN_BYTES = 1 << 20 +_OWNER_REPLICA_HEADER = "x-openshell-owner" +_OWNER_REQUEST_HEADER = "x-openshell-want-owner" +_ROUTE_REPLICA_HEADER = "x-openshell-replica" T = TypeVar("T") @@ -99,6 +102,50 @@ def all(self) -> builtins.list[T]: return [item for page in self for item in page.items] +class _OwnerCache: + """Owner replicas the gateway last reported, keyed by workspace and name. + + An empty owner means the gateway reported none. A stale owner is safe: a + live non-owner replica relays the call. + """ + + _TTL_SECONDS = 30.0 + _MAX_ENTRIES = 1024 + + def __init__(self) -> None: + self._lock = threading.Lock() + self._entries: dict[tuple[str, str], tuple[float, str]] = {} + + def get(self, workspace: str, name: str) -> str | None: + with self._lock: + entry = self._entries.get((workspace, name)) + if entry is None or time.monotonic() - entry[0] >= self._TTL_SECONDS: + return None + return entry[1] + + def put(self, workspace: str, name: str, owner: str) -> None: + with self._lock: + if len(self._entries) >= self._MAX_ENTRIES: + self._entries.clear() + self._entries[(workspace, name)] = (time.monotonic(), owner) + + def forget(self, workspace: str, name: str) -> None: + with self._lock: + self._entries.pop((workspace, name), None) + + +def _retry_unrouted(error: grpc.RpcError) -> bool: + """Envoy Gateway answers a replica route with no ready endpoints with + UNAVAILABLE and no message; gateway errors always carry one.""" + status = from_grpc_error(error) + details = status.details() or "" + return status.code() == grpc.StatusCode.UNAVAILABLE and ( + details == "" + or "no healthy upstream" in details + or "connection failure" in details + ) + + def _workspace_scope(workspace: str) -> datamodel_pb2.WorkspaceSelector: if not workspace: raise ValueError("workspace must be non-empty") @@ -730,6 +777,7 @@ def __init__( self._stub = openshell_pb2_grpc.OpenShellStub( _error_mapping_channel(self._channel) ) + self._owners = _OwnerCache() @classmethod def from_active_cluster( @@ -1280,7 +1328,7 @@ def exec_stream( 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) + stream = self._open_exec(request, workspace, grpc_deadline) stdout_parts: list[bytes] = [] stderr_parts: list[bytes] = [] @@ -1308,6 +1356,53 @@ def exec_stream( stderr=b"".join(stderr_parts).decode("utf-8", errors="replace"), ) + def _owner_of(self, workspace: str, sandbox: str) -> str: + """Best effort: a failed lookup leaves the call unrouted.""" + owner = self._owners.get(workspace, sandbox) + if owner is not None: + return owner + try: + _, call = self._stub.GetSandbox.with_call( + openshell_pb2.GetSandboxRequest( + workspace_scope=_workspace_scope(workspace), + name=sandbox, + ), + timeout=self._timeout, + metadata=((_OWNER_REQUEST_HEADER, "1"),), + ) + except grpc.RpcError: + return "" + owner = dict(call.initial_metadata() or ()).get(_OWNER_REPLICA_HEADER, "") + self._owners.put(workspace, sandbox, owner) + return owner + + def _open_exec( + self, + request: openshell_pb2.ExecSandboxRequest, + workspace: str, + timeout: float, + ) -> Iterator[openshell_pb2.ExecSandboxEvent]: + """Open ExecSandbox on the owner replica, retrying once unrouted.""" + owner = self._owner_of(workspace, request.sandbox) + if owner: + events = self._stub.ExecSandbox( + request, + timeout=timeout, + metadata=((_ROUTE_REPLICA_HEADER, owner),), + ) + try: + first = next(events, 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 + yield from events + return + yield from self._stub.ExecSandbox(request, timeout=timeout) + def exec( self, sandbox: str, diff --git a/python/openshell/sandbox_test.py b/python/openshell/sandbox_test.py index c1df0d46e0..f35a9e4c6a 100644 --- a/python/openshell/sandbox_test.py +++ b/python/openshell/sandbox_test.py @@ -3,22 +3,25 @@ from __future__ import annotations +import contextlib import json import os import pickle import threading import time +from concurrent import futures from copy import deepcopy from dataclasses import asdict from pathlib import Path from types import SimpleNamespace -from typing import Any, cast +from typing import TYPE_CHECKING, Any, Never, cast +import grpc import pytest from google.protobuf.struct_pb2 import Struct import openshell.sandbox as sandbox_module -from openshell._proto import openshell_pb2 +from openshell._proto import openshell_pb2, openshell_pb2_grpc from openshell.mutations import DeletionOutcome from openshell.sandbox import ( _OIDC_TOKEN_EXPIRY_GRACE_SECONDS, @@ -43,11 +46,15 @@ _make_cluster_bearer_provider, _normalize_bearer, _OidcRefresher, + _OwnerCache, _read_oidc_token_bundle, _sandbox_ref, _validate_oauth_url, ) +if TYPE_CHECKING: + from collections.abc import Iterator + def _request_workspace(request: Any) -> str | None: scope = request.workspace_scope @@ -412,6 +419,11 @@ def handler(request: Any) -> Any: class _FakeStub: + class GetSandbox: + @staticmethod + def with_call(*_args: object, **_kwargs: object) -> Never: + raise grpc.RpcError() + def __init__(self) -> None: self.request: openshell_pb2.ExecSandboxRequest | None = None @@ -431,9 +443,102 @@ def _client_with_fake_stub(stub: object) -> SandboxClient: client = object.__new__(SandboxClient) client._timeout = 30.0 client._stub = cast("Any", stub) + client._owners = _OwnerCache() return client +class _RoutingGateway(openshell_pb2_grpc.OpenShellServicer): + """Gateway that reports an owner and, behind a replica route, can answer + like Envoy with a trailers-only UNAVAILABLE.""" + + def __init__( + self, owner: str, *, reject_routed: str | None = None, fail: str = "" + ) -> None: + self.owner = owner + self.reject_routed = reject_routed + self.fail = fail + self.wanted_owner: list[str] = [] + self.routes: list[str] = [] + + def GetSandbox(self, request, context): + _ = request + headers = dict(context.invocation_metadata()) + self.wanted_owner.append(headers.get("x-openshell-want-owner", "")) + if self.owner: + context.send_initial_metadata((("x-openshell-owner", self.owner),)) + return openshell_pb2.SandboxResponse() + + def ExecSandbox(self, request, context): + _ = request + route = dict(context.invocation_metadata()).get("x-openshell-replica", "") + self.routes.append(route) + if route and self.reject_routed is not None: + context.abort(grpc.StatusCode.UNAVAILABLE, self.reject_routed) + if self.fail: + context.abort(grpc.StatusCode.UNAVAILABLE, self.fail) + yield openshell_pb2.ExecSandboxEvent( + stdout=openshell_pb2.ExecSandboxStdout(data=b"ok") + ) + yield openshell_pb2.ExecSandboxEvent( + exit=openshell_pb2.ExecSandboxExit(exit_code=0) + ) + + +@contextlib.contextmanager +def _routing_client(gateway: _RoutingGateway) -> Iterator[SandboxClient]: + server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) + openshell_pb2_grpc.add_OpenShellServicer_to_server(gateway, server) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + try: + with SandboxClient(f"127.0.0.1:{port}", timeout=5) as client: + yield client + finally: + server.stop(None) + + +def test_exec_routes_to_reported_owner_and_caches_it() -> None: + gateway = _RoutingGateway("openshell-1") + with _routing_client(gateway) as client: + for _ in range(2): + result = client.exec("box", ["true"], workspace="default") + assert result.stdout == "ok" + + assert gateway.wanted_owner == ["1"] + assert gateway.routes == ["openshell-1", "openshell-1"] + + +def test_exec_without_owner_sends_no_hint() -> None: + gateway = _RoutingGateway("") + with _routing_client(gateway) as client: + assert client.exec("box", ["true"], workspace="default").exit_code == 0 + + assert gateway.routes == [""] + + +@pytest.mark.parametrize("message", ["", "no healthy upstream"]) +def test_exec_retries_once_unrouted_when_owner_route_is_empty(message: str) -> None: + gateway = _RoutingGateway("openshell-1", reject_routed=message) + with _routing_client(gateway) as client: + for _ in range(2): + result = client.exec("box", ["true"], workspace="default") + assert result.stdout == "ok" + + assert gateway.routes == ["openshell-1", "", "openshell-1", ""] + assert gateway.wanted_owner == ["1", "1"] + + +def test_exec_does_not_retry_gateway_errors() -> None: + gateway = _RoutingGateway("openshell-1", fail="supervisor not connected") + with ( + _routing_client(gateway) as client, + pytest.raises(grpc.RpcError, match="supervisor not connected"), + ): + client.exec("box", ["true"], workspace="default") + + assert gateway.routes == ["openshell-1"] + + def _template_client_with_fake_stub(stub: object) -> SandboxTemplateClient: client = object.__new__(SandboxTemplateClient) client._timeout = 30.0 From a33b712edf3dd19dce62b4c27ee480f7a220dfde Mon Sep 17 00:00:00 2001 From: divesh Date: Fri, 30 Oct 2026 18:00:54 -0700 Subject: [PATCH 2/2] refactor(sdk): share the Python GetSandbox request for owner lookup Reuse get's request for the owner lookup, point the routing headers at their documented contract, and accept only string owner metadata. Part of #4223 Signed-off-by: divesh --- python/openshell/sandbox.py | 20 +++++++++++--------- python/openshell/sandbox_test.py | 4 ++-- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index d9ac242f23..c02a33cdbf 100644 --- a/python/openshell/sandbox.py +++ b/python/openshell/sandbox.py @@ -40,6 +40,7 @@ _OAUTH_MAX_RESPONSE_BYTES = 1 << 20 _PAGER_MAX_CONSUMED_TOKENS = 10_000 _PAGER_MAX_CONSUMED_TOKEN_BYTES = 1 << 20 +# Replica routing headers and retry rule: docs/kubernetes/ingress.mdx. _OWNER_REPLICA_HEADER = "x-openshell-owner" _OWNER_REQUEST_HEADER = "x-openshell-want-owner" _ROUTE_REPLICA_HEADER = "x-openshell-replica" @@ -152,6 +153,12 @@ def _workspace_scope(workspace: str) -> datamodel_pb2.WorkspaceSelector: return datamodel_pb2.WorkspaceSelector(workspace=workspace) +def _get_sandbox_request(name: str, workspace: str) -> openshell_pb2.GetSandboxRequest: + return openshell_pb2.GetSandboxRequest( + workspace_scope=_workspace_scope(workspace), name=name + ) + + def _all_workspaces_scope() -> datamodel_pb2.WorkspaceSelector: return datamodel_pb2.WorkspaceSelector(all_workspaces=datamodel_pb2.AllWorkspaces()) @@ -1015,10 +1022,7 @@ def providers(self) -> ProviderClient: def get(self, name: str, *, workspace: str) -> SandboxRef: response = self._stub.GetSandbox( - openshell_pb2.GetSandboxRequest( - workspace_scope=_workspace_scope(workspace), - name=name, - ), + _get_sandbox_request(name, workspace), timeout=self._timeout, ) return _sandbox_ref(response.sandbox) @@ -1363,16 +1367,14 @@ def _owner_of(self, workspace: str, sandbox: str) -> str: return owner try: _, call = self._stub.GetSandbox.with_call( - openshell_pb2.GetSandboxRequest( - workspace_scope=_workspace_scope(workspace), - name=sandbox, - ), + _get_sandbox_request(sandbox, workspace), timeout=self._timeout, metadata=((_OWNER_REQUEST_HEADER, "1"),), ) except grpc.RpcError: return "" - owner = dict(call.initial_metadata() or ()).get(_OWNER_REPLICA_HEADER, "") + value = dict(call.initial_metadata() or ()).get(_OWNER_REPLICA_HEADER) + owner = value if isinstance(value, str) else "" self._owners.put(workspace, sandbox, owner) return owner diff --git a/python/openshell/sandbox_test.py b/python/openshell/sandbox_test.py index f35a9e4c6a..3fed31734e 100644 --- a/python/openshell/sandbox_test.py +++ b/python/openshell/sandbox_test.py @@ -53,7 +53,7 @@ ) if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Generator def _request_workspace(request: Any) -> str | None: @@ -485,7 +485,7 @@ def ExecSandbox(self, request, context): @contextlib.contextmanager -def _routing_client(gateway: _RoutingGateway) -> Iterator[SandboxClient]: +def _routing_client(gateway: _RoutingGateway) -> Generator[SandboxClient]: server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) openshell_pb2_grpc.add_OpenShellServicer_to_server(gateway, server) port = server.add_insecure_port("127.0.0.1:0")