sandbox.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import base64
import contextlib
import errno
import ipaddress
import json
import math
import os
import pathlib
import sys
import tempfile
import threading
import time
from collections import namedtuple
from dataclasses import dataclass, field
from enum import IntEnum
from typing import TYPE_CHECKING, Any, Generic, Never, SupportsIndex, TypeVar, cast
from urllib.parse import urlparse

import grpc
import httpx

from ._proto import (
    datamodel_pb2,
    openshell_pb2,
    openshell_pb2_grpc,
)
from .errors import _error_mapping_channel, from_grpc_error
from .mutations import DeletionOutcome, DeletionResult

_ClientCallDetailsBase = namedtuple(
    "_ClientCallDetailsBase",
    ("method", "timeout", "metadata", "credentials", "wait_for_ready", "compression"),
)

_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"
T = TypeVar("T")


@dataclass(frozen=True)
class Page(Generic[T]):
    """One response page from a list operation."""

    items: builtins.list[T]
    next_page_token: str


class Pager(Generic[T]):
    """Lazy, single-pass iterator over the continuation-token contract.

    The repeated-token guard has bounded memory and raises ``SandboxError`` if
    the traversal exceeds that guard's token-count or byte budget.
    """

    def __init__(self, fetch: Callable[[str], Page[T]], page_token: str = "") -> None:
        self._fetch = fetch
        self._page_token: str | None = page_token
        self._consumed_page_tokens: set[str] = set()
        self._consumed_page_token_bytes = 0

    def __iter__(self) -> Pager[T]:
        return self

    def _validate_page_token_budget(self, page_token: str) -> int:
        if not page_token:
            return 0
        token_bytes = len(page_token.encode("utf-8"))
        if (
            len(self._consumed_page_tokens) >= _PAGER_MAX_CONSUMED_TOKENS
            or self._consumed_page_token_bytes + token_bytes
            > _PAGER_MAX_CONSUMED_TOKEN_BYTES
        ):
            raise SandboxError("pager continuation token history limit exceeded")
        return token_bytes

    def __next__(self) -> Page[T]:
        if self._page_token is None:
            raise StopIteration
        page_token = self._page_token
        token_bytes = self._validate_page_token_budget(page_token)
        page = self._fetch(page_token)
        if page_token:
            self._consumed_page_tokens.add(page_token)
            self._consumed_page_token_bytes += token_bytes
        next_page_token = page.next_page_token
        if next_page_token and next_page_token in self._consumed_page_tokens:
            raise SandboxError("pager received a repeated continuation token")
        self._page_token = next_page_token or None
        return page

    def all(self) -> builtins.list[T]:
        """Consume the pager and collect every remaining item."""
        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")
    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())


def _service_exposure_messages(
    exposures: Sequence[ServiceExposure] | None,
) -> list[openshell_pb2.SandboxServiceExposure]:
    return [
        openshell_pb2.SandboxServiceExposure(
            service=exposure.service,
            target_port=exposure.target_port,
            authorization_mode=(
                openshell_pb2.SERVICE_AUTHORIZATION_MODE_BEARER_PASSTHROUGH
                if exposure.authorization_mode
                == ServiceAuthorizationMode.BEARER_PASSTHROUGH
                else openshell_pb2.SERVICE_AUTHORIZATION_MODE_STRIP
            ),
        )
        for exposure in exposures or ()
    ]


class _ClientCallDetails(_ClientCallDetailsBase, grpc.ClientCallDetails):
    pass


if TYPE_CHECKING:
    import builtins
    from collections.abc import Callable, Iterator, Mapping, Sequence

    from .providers import ProviderClient


@dataclass(frozen=True)
class TlsConfig:
    """Channel TLS material.

    All three fields are optional so callers can pick the trust profile:

    - Full mTLS: pass all three (server trusts client identity).
    - CA-only: pass `ca_path` (custom CA, no client identity).
    - System roots: pass no fields (`TlsConfig()`) — uses the OS trust
      store. Useful for OIDC gateways behind a public CA.

    `cert_path` and `key_path` must be set together or not at all.
    """

    ca_path: pathlib.Path | None = None
    cert_path: pathlib.Path | None = None
    key_path: pathlib.Path | None = None

    def __post_init__(self) -> None:
        if (self.cert_path is None) != (self.key_path is None):
            raise ValueError("TlsConfig: cert_path and key_path must be set together")


class _BearerAuthInterceptor(
    grpc.UnaryUnaryClientInterceptor,
    grpc.UnaryStreamClientInterceptor,
    grpc.StreamUnaryClientInterceptor,
    grpc.StreamStreamClientInterceptor,
):
    """Add `authorization: Bearer <token>` to every outgoing RPC.

    Implemented as an interceptor (not call credentials) so it works on
    both plaintext and TLS channels without needing
    `grpc.composite_channel_credentials`. The token provider is invoked
    per call, so callers can swap tokens at runtime by mutating shared
    state or returning a fresh value from the callable.
    """

    def __init__(self, token_provider: Callable[[], str]) -> None:
        self._token_provider = token_provider

    def _attach(self, details: grpc.ClientCallDetails) -> grpc.ClientCallDetails:
        original_metadata = getattr(details, "metadata", None)
        metadata = list(original_metadata) if original_metadata else []
        metadata.append(("authorization", f"Bearer {self._token_provider()}"))
        return _ClientCallDetails(
            getattr(details, "method", None),
            getattr(details, "timeout", None),
            metadata,
            getattr(details, "credentials", None),
            getattr(details, "wait_for_ready", None),
            getattr(details, "compression", None),
        )

    def intercept_unary_unary(self, continuation, client_call_details, request):
        return continuation(self._attach(client_call_details), request)

    def intercept_unary_stream(self, continuation, client_call_details, request):
        return continuation(self._attach(client_call_details), request)

    def intercept_stream_unary(
        self, continuation, client_call_details, request_iterator
    ):
        return continuation(self._attach(client_call_details), request_iterator)

    def intercept_stream_stream(
        self, continuation, client_call_details, request_iterator
    ):
        return continuation(self._attach(client_call_details), request_iterator)


def _normalize_bearer(
    bearer: str | Callable[[], str] | None,
) -> Callable[[], str] | None:
    if bearer is None:
        return None
    if callable(bearer):
        return cast("Callable[[], str]", bearer)
    token = bearer
    return lambda: token


def _is_loopback_host(hostname: str | None) -> bool:
    if hostname is None:
        return False
    if hostname.lower() == "localhost":
        return True
    with contextlib.suppress(ValueError):
        return ipaddress.ip_address(hostname).is_loopback
    return False


def _is_local_grpc_endpoint(endpoint: str) -> bool:
    if endpoint.startswith("unix:"):
        return True
    return _is_loopback_host(urlparse(f"//{endpoint}").hostname)


def _validate_oauth_url(name: str, raw: str) -> str:
    """Validate an OAuth endpoint without reflecting attacker-controlled URLs."""
    parsed = urlparse(raw)
    if not parsed.scheme or not parsed.hostname or parsed.username or parsed.password:
        raise SandboxError(f"invalid OAuth {name} URL")
    if parsed.fragment:
        raise SandboxError(f"OAuth {name} URL must not contain a fragment")
    if parsed.scheme == "https":
        return raw
    if parsed.scheme == "http" and _is_loopback_host(parsed.hostname):
        return raw
    raise SandboxError(
        f"OAuth {name} URL must use HTTPS (HTTP is allowed only for loopback hosts)"
    )


def _oauth_json_request(
    client: httpx.Client,
    method: str,
    url: str,
    *,
    kind: str,
    headers: Mapping[str, str],
    data: Mapping[str, str] | None = None,
) -> tuple[int, dict[str, object]]:
    """Read a bounded JSON response without retaining arbitrary error bodies."""
    with client.stream(method, url, headers=headers, data=data) as response:
        if response.status_code != 200:
            return response.status_code, {}
        content = bytearray()
        for chunk in response.iter_bytes():
            content.extend(chunk)
            if len(content) > _OAUTH_MAX_RESPONSE_BYTES:
                raise SandboxError(f"OAuth {kind} response is too large")
    try:
        payload = json.loads(content)
    except (UnicodeDecodeError, json.JSONDecodeError):
        raise SandboxError(f"OAuth {kind} returned invalid JSON") from None
    if not isinstance(payload, dict):
        raise SandboxError(f"OAuth {kind} returned invalid JSON")
    return 200, payload


class ClientCredentialsAuth:
    """Renewable OAuth 2.0 client-credentials bearer provider.

    Tokens and client secrets remain in memory. The first RPC lazily performs
    discovery and an exchange; later RPCs share the cached token until it is
    within 30 seconds of expiry. Concurrent callers share one exchange.
    """

    def __init__(
        self,
        *,
        client_secret: str | Callable[[], str],
        issuer: str | None = None,
        client_id: str | None = None,
        scopes: Sequence[str] | None = None,
        audience: str | None = None,
        insecure: bool = False,
        timeout: float = 30.0,
        _transport: httpx.BaseTransport | None = None,
    ) -> None:
        if not isinstance(client_secret, str) and not callable(client_secret):
            raise TypeError("client_secret must be a string or zero-argument callable")
        if isinstance(client_secret, str) and not client_secret:
            raise SandboxError("OAuth client secret must not be empty")
        self._secret = client_secret
        self._issuer = issuer
        self._client_id = client_id
        self._scopes = tuple(scopes) if scopes is not None else None
        self._audience = audience
        self._insecure = insecure
        self._timeout = timeout
        self._transport = _transport
        self._lock = threading.Lock()
        self._access_token: str | None = None
        self._expires_at = 0.0
        self._token_endpoint: str | None = None

    def __repr__(self) -> str:
        return (
            "ClientCredentialsAuth(issuer="
            f"{self._issuer!r}, client_id={self._client_id!r}, "
            f"scopes={self._scopes!r}, audience={self._audience!r})"
        )

    def _apply_gateway_metadata(
        self, metadata: Mapping[str, object], *, insecure: bool = False
    ) -> None:
        """Fill omitted fields and apply active-gateway transport settings."""
        if self._issuer is None:
            value = metadata.get("oidc_issuer")
            self._issuer = value if isinstance(value, str) and value else None
        if self._client_id is None:
            value = metadata.get("oidc_client_id")
            self._client_id = value if isinstance(value, str) and value else None
        if self._audience is None:
            value = metadata.get("oidc_audience")
            self._audience = value if isinstance(value, str) and value else None
        if self._scopes is None:
            value = metadata.get("oidc_scopes")
            if isinstance(value, str):
                self._scopes = tuple(value.split())
        # Preserve an explicit provider-level opt-in while also honoring the
        # active-gateway construction flag used by the existing OIDC refresher.
        self._insecure = self._insecure or insecure

    def __call__(self) -> str:
        if (
            self._access_token is not None
            and time.time() + _OIDC_TOKEN_EXPIRY_GRACE_SECONDS < self._expires_at
        ):
            return self._access_token
        with self._lock:
            if (
                self._access_token is not None
                and time.time() + _OIDC_TOKEN_EXPIRY_GRACE_SECONDS < self._expires_at
            ):
                return self._access_token
            token, expires_at = self._exchange()
            self._access_token = token
            self._expires_at = expires_at
            return token

    def _exchange(self) -> tuple[str, float]:
        if not self._issuer or not self._client_id:
            raise SandboxError("OAuth client credentials require issuer and client_id")
        issuer = _validate_oauth_url("issuer", self._issuer).rstrip("/")
        headers = {"accept": "application/json"}
        try:
            with httpx.Client(
                verify=not self._insecure,
                follow_redirects=False,
                timeout=self._timeout,
                transport=self._transport,
            ) as client:
                if self._token_endpoint is None:
                    status, discovery = _oauth_json_request(
                        client,
                        "GET",
                        f"{issuer}/.well-known/openid-configuration",
                        kind="discovery",
                        headers=headers,
                    )
                    if status != 200:
                        raise SandboxError(f"OAuth discovery failed with HTTP {status}")
                    discovered = discovery.get("issuer")
                    if (
                        not isinstance(discovered, str)
                        or discovered.rstrip("/") != issuer
                    ):
                        raise SandboxError("OAuth discovery issuer mismatch")
                    endpoint = discovery.get("token_endpoint")
                    if not isinstance(endpoint, str) or not endpoint:
                        raise SandboxError(
                            "OAuth discovery document is missing token_endpoint"
                        )
                    self._token_endpoint = _validate_oauth_url(
                        "token endpoint", endpoint
                    )
                secret = self._resolve_secret()
                data = {
                    "grant_type": "client_credentials",
                    "client_id": self._client_id,
                    "client_secret": secret,
                }
                if self._scopes:
                    data["scope"] = " ".join(self._scopes)
                if self._audience:
                    data["audience"] = self._audience
                status, payload = _oauth_json_request(
                    client,
                    "POST",
                    self._token_endpoint,
                    kind="client credentials response",
                    data=data,
                    headers=headers,
                )
        except SandboxError:
            raise
        except httpx.HTTPError:
            raise SandboxError("OAuth client credentials request failed") from None
        if status != 200:
            raise SandboxError(
                f"OAuth client credentials exchange failed with HTTP {status}"
            )
        access_token = payload.get("access_token")
        expires_in = payload.get("expires_in")
        if not isinstance(access_token, str) or not access_token:
            raise SandboxError(
                "OAuth client credentials response is missing access_token"
            )
        if (
            isinstance(expires_in, bool)
            or not isinstance(expires_in, (int, float))
            or not math.isfinite(expires_in)
            or expires_in <= 0
        ):
            raise SandboxError(
                "OAuth client credentials response requires a positive finite expires_in"
            )
        return access_token, time.time() + float(expires_in)

    def _resolve_secret(self) -> str:
        if isinstance(self._secret, str):
            return self._secret
        try:
            value = self._secret()
        except Exception:
            raise SandboxError("OAuth client secret supplier failed") from None
        if not isinstance(value, str) or not value:
            raise SandboxError("OAuth client secret supplier returned an empty secret")
        return value


@dataclass(frozen=True)
class SandboxConditionRef:
    type: str
    status: str
    reason: str = ""
    message: str = ""
    transition_time_ms: int | None = None


@dataclass(frozen=True)
class SandboxStatusRef:
    phase: int
    current_policy_version: int
    exit_code: int | None = None
    restart_count: int = 0
    next_restart_at_ms: int | None = None
    main_process_started_at_ms: int | None = None
    conditions: tuple[SandboxConditionRef, ...] = ()


class ServiceAuthorizationMode(IntEnum):
    """Handling for an incoming application Authorization header."""

    STRIP = 1
    BEARER_PASSTHROUGH = 2


@dataclass(frozen=True)
class ServiceExposure:
    """A loopback HTTP service to expose during sandbox creation."""

    target_port: int
    service: str = ""
    authorization_mode: ServiceAuthorizationMode = ServiceAuthorizationMode.STRIP


class _ImmutableLabels(dict[str, str]):
    """A read-only, copy- and pickle-safe label mapping."""

    def _deny_mutation(self, *args: object, **kwargs: object) -> Never:
        del args, kwargs
        raise TypeError("sandbox labels are immutable")

    __setitem__ = _deny_mutation
    __delitem__ = _deny_mutation
    clear = _deny_mutation
    pop = _deny_mutation
    popitem = _deny_mutation
    setdefault = _deny_mutation
    update = _deny_mutation
    __ior__ = _deny_mutation

    def __deepcopy__(self, memo: dict[int, object]) -> _ImmutableLabels:
        del memo
        return type(self)(self)

    def __reduce_ex__(self, protocol: SupportsIndex, /) -> str | tuple[Any, ...]:
        del protocol
        return type(self), (dict(self),)


@dataclass(frozen=True)
class SandboxWorkloadTemplateProvenanceRef:
    name: str
    resource_version: str


@dataclass(frozen=True)
class SandboxRef:
    id: str
    name: str
    workspace: str
    status: SandboxStatusRef
    # Excluded from equality/hash to preserve the original identity while the
    # immutable mapping remains safe for deepcopy, pickle, and asdict.
    labels: Mapping[str, str] = field(default_factory=_ImmutableLabels, compare=False)
    created_from_workload_template: SandboxWorkloadTemplateProvenanceRef | None = None
    # Populated by create operations. The empty key identifies the unnamed service.
    service_urls: Mapping[str, str] = field(
        default_factory=_ImmutableLabels, compare=False
    )
    # Stable SSH host identity; absent on older runtimes.
    host_key_fingerprint: str | None = None

    def __post_init__(self) -> None:
        object.__setattr__(self, "labels", _ImmutableLabels(self.labels))
        object.__setattr__(self, "service_urls", _ImmutableLabels(self.service_urls))

    @property
    def phase(self) -> int:
        return self.status.phase

    @property
    def current_policy_version(self) -> int:
        return self.status.current_policy_version

    @property
    def conditions(self) -> tuple[SandboxConditionRef, ...]:
        return self.status.conditions


@dataclass(frozen=True)
class ExecChunk:
    stream: str
    data: bytes


@dataclass(frozen=True)
class ExecResult:
    exit_code: int
    stdout: str
    stderr: str


class SandboxError(RuntimeError):
    pass


class SandboxSession:
    def __init__(self, client: SandboxClient, sandbox: SandboxRef) -> None:
        self._client = client
        self.sandbox = sandbox
        self._workspace = sandbox.workspace

    @property
    def id(self) -> str:
        return self.sandbox.id

    def exec(
        self,
        command: Sequence[str],
        *,
        stream_output: bool = False,
        workdir: str | None = None,
        env: Mapping[str, str] | None = None,
        stdin: bytes | None = None,
        timeout_seconds: int | None = None,
        no_login_shell: bool = False,
    ) -> ExecResult:
        return self._client.exec(
            self.sandbox.name,
            command,
            workspace=self._workspace,
            stream_output=stream_output,
            workdir=workdir,
            env=env,
            stdin=stdin,
            timeout_seconds=timeout_seconds,
            no_login_shell=no_login_shell,
        )

    def exec_python(
        self,
        function: Callable[..., object],
        *,
        args: Sequence[object] = (),
        kwargs: Mapping[str, object] | None = None,
        stream_output: bool = False,
        workdir: str | None = None,
        env: Mapping[str, str] | None = None,
        timeout_seconds: int | None = None,
    ) -> ExecResult:
        return self._client.exec_python(
            self.sandbox.name,
            function,
            workspace=self._workspace,
            args=args,
            kwargs=kwargs,
            stream_output=stream_output,
            workdir=workdir,
            env=env,
            timeout_seconds=timeout_seconds,
        )

    def delete(self, *, allow_missing: bool = False) -> DeletionResult:
        return self._client.delete(
            self.sandbox.name, workspace=self._workspace, allow_missing=allow_missing
        )

    def stop(self) -> SandboxRef:
        self.sandbox = self._client.stop(self.sandbox.name, workspace=self._workspace)
        return self.sandbox

    def start(self) -> SandboxRef:
        self.sandbox = self._client.start(self.sandbox.name, workspace=self._workspace)
        return self.sandbox

    def suspend(self) -> SandboxRef:
        self.sandbox = self._client.suspend(
            self.sandbox.name, workspace=self._workspace
        )
        return self.sandbox

    def resume(self, *, restart_if_restore_fails: bool = False) -> SandboxRef:
        self.sandbox = self._client.resume(
            self.sandbox.name,
            workspace=self._workspace,
            restart_if_restore_fails=restart_if_restore_fails,
        )
        return self.sandbox


class SandboxClient:
    """gRPC client for sandbox CRUD and command execution."""

    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:
        """Create a SandboxClient.

        Args:
            endpoint: host:port for the gateway gRPC service.
            tls: mTLS material. None for a plaintext channel.
            bearer_token: OIDC access token, or a zero-arg callable
                returning the current token (called per RPC; supports
                runtime refresh). Combines with `tls` — pass both when
                the gateway uses mTLS for transport identity and OIDC
                for user identity.
            client_credentials: renewable OAuth client-credentials provider.
                Mutually exclusive with `bearer_token`. A non-loopback endpoint
                requires `tls` so the acquired bearer is never sent in cleartext.
            timeout: default per-call timeout in seconds.
            cluster_name: optional friendly name for error messages.
            _bearer_close: internal — wired by `from_active_cluster`
                when an `_OidcRefresher` owns the bearer callable, so
                `close()` can release the refresher's HTTP client.
                Public callers should not pass this; they own the
                lifecycle of any callable they supplied as
                `bearer_token`.
        """
        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._endpoint = endpoint
        self._timeout = timeout
        self._cluster_name = cluster_name
        self._bearer_close = _bearer_close
        if tls is None:
            self._channel = grpc.insecure_channel(endpoint)
        else:
            # Build credentials from whatever subset of mTLS material the
            # caller supplied. None for `root_certificates` makes gRPC use
            # the system trust store, which is what we want for OIDC
            # gateways behind a public CA.
            credentials = grpc.ssl_channel_credentials(
                root_certificates=(tls.ca_path.read_bytes() if tls.ca_path else None),
                private_key=(tls.key_path.read_bytes() if tls.key_path else None),
                certificate_chain=(
                    tls.cert_path.read_bytes() if tls.cert_path else None
                ),
            )
            self._channel = grpc.secure_channel(endpoint, credentials)
        provider = _normalize_bearer(client_credentials or bearer_token)
        if provider is not None:
            self._channel = grpc.intercept_channel(
                self._channel,
                _BearerAuthInterceptor(provider),
            )
        self._stub = openshell_pb2_grpc.OpenShellStub(
            _error_mapping_channel(self._channel)
        )
        self._owners = _OwnerCache()

    @classmethod
    def from_active_cluster(
        cls,
        *,
        cluster: str | None = None,
        timeout: float = 30.0,
        auto_refresh: bool = True,
        write_back: bool = True,
        insecure: bool = False,
        client_credentials: ClientCredentialsAuth | None = None,
    ) -> SandboxClient:
        """Construct a `SandboxClient` from the active gateway's on-disk state.

        Args:
            cluster: explicit gateway name; otherwise reads
                `$OPENSHELL_GATEWAY` or `~/.config/openshell/active_gateway`.
            timeout: per-call gRPC timeout in seconds.
            auto_refresh: when True (default) and the gateway uses OIDC,
                lazily refresh the access token via the IdP's token endpoint
                if the cached `oidc_token.json` is near expiry. Matches the
                lazy-refresh patterns used by `google-auth` and `botocore`.
                Set False to keep the SDK as a read-only consumer of the
                CLI's cache (fail closed on expiry).
            write_back: when True (default, and `auto_refresh=True`),
                atomically persist refreshed bundles back to
                `oidc_token.json` so other processes — including the
                Rust CLI — see the rotation. Required for IdPs with
                refresh-token rotation enabled (Keycloak, Entra in
                strict mode): an in-memory-only refresh would leave the
                on-disk `refresh_token` pointing at an invalidated
                value, and any other process starting from that disk
                state would fail on its first refresh. Set False only
                when you know the SDK is the sole consumer of this
                gateway directory.
            insecure: when True, disables TLS certificate verification
                for OIDC discovery and refresh calls. Mirrors the Rust
                CLI's `--insecure` flag for issuers behind self-signed
                certs. Off by default.
            client_credentials: renewable OAuth client-credentials provider.
                Omitted issuer, client ID, audience, and scopes are filled from
                the registered gateway metadata. This provider does not read or
                write `oidc_token.json`. Remote plaintext gateways are rejected.
        """
        cluster_name = cluster or _resolve_active_cluster()
        gateway_dir = _xdg_config_home() / "openshell" / "gateways" / cluster_name
        metadata_path = gateway_dir / "metadata.json"
        try:
            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
        except FileNotFoundError:
            raise SandboxError(f"gateway '{cluster_name}' not found") from None
        if "gateway_endpoint" not in metadata:
            raise SandboxError(f"gateway '{cluster_name}' metadata missing endpoint")
        parsed = urlparse(metadata["gateway_endpoint"])
        host = parsed.hostname or "127.0.0.1"
        port = parsed.port or (443 if parsed.scheme == "https" else 80)
        endpoint = f"{host}:{port}"

        # TLS transport. Mirror crates/openshell-tui/src/lib.rs
        # `build_oidc_channel` — for an https gateway, always build a
        # secure channel and pick the strongest available trust profile.
        tls: TlsConfig | None = None
        if parsed.scheme == "https":
            mtls_dir = gateway_dir / "mtls"
            ca = mtls_dir / "ca.crt" if (mtls_dir / "ca.crt").exists() else None
            cert = mtls_dir / "tls.crt" if (mtls_dir / "tls.crt").exists() else None
            key = mtls_dir / "tls.key" if (mtls_dir / "tls.key").exists() else None
            if ca is not None and cert is not None and key is not None:
                # Full mTLS.
                tls = TlsConfig(ca_path=ca, cert_path=cert, key_path=key)
            elif ca is not None:
                # CA-only trust (no client identity).
                tls = TlsConfig(ca_path=ca)
            else:
                # System roots (e.g. OIDC gateway behind a public CA).
                tls = TlsConfig()

        # OIDC bearer. Mirror the Rust CLI/TUI: the gateway metadata's
        # `auth_mode` is authoritative — a stale oidc_token.json next to
        # a non-OIDC gateway should NOT cause us to attach a bearer.
        bearer_token: Callable[[], str] | None = None
        bearer_close: Callable[[], None] | None = None
        if client_credentials is not None:
            if metadata.get("auth_mode") != "oidc":
                raise SandboxError(
                    f"gateway '{cluster_name}' is not configured for OIDC"
                )
            client_credentials._apply_gateway_metadata(metadata, insecure=insecure)
        elif metadata.get("auth_mode") == "oidc":
            bearer_token, bearer_close = _make_cluster_bearer_provider(
                gateway_dir,
                cluster_name,
                auto_refresh=auto_refresh,
                write_back=write_back,
                insecure=insecure,
            )

        return cls(
            endpoint,
            tls=tls,
            bearer_token=bearer_token,
            client_credentials=client_credentials,
            timeout=timeout,
            cluster_name=cluster_name,
            _bearer_close=bearer_close,
        )

    def close(self) -> None:
        """Release the gRPC channel and any bearer-auth resources.

        Idempotent. If `from_active_cluster` wired up an OIDC refresher
        for this client, the refresher's underlying httpx.Client is
        closed here too — otherwise long-lived services that churn
        clients would leak sockets / file descriptors until GC.
        """
        self._channel.close()
        if self._bearer_close is not None:
            with contextlib.suppress(Exception):
                self._bearer_close()
            self._bearer_close = None

    def __enter__(self) -> SandboxClient:
        return self

    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 create(
        self,
        *,
        workspace: str,
        spec: openshell_pb2.SandboxSpec | None = None,
        name: str | None = None,
        labels: Mapping[str, str] | None = None,
        service_exposures: Sequence[ServiceExposure] | None = None,
    ) -> SandboxRef:
        request_spec = spec if spec is not None else _default_spec()
        response = self._stub.CreateSandbox(
            openshell_pb2.CreateSandboxRequest(
                workspace_scope=_workspace_scope(workspace),
                spec=request_spec,
                name=name or "",
                labels=dict(labels) if labels else {},
                service_exposures=_service_exposure_messages(service_exposures),
            ),
            timeout=self._timeout,
        )
        sandbox_ref = _sandbox_ref(response.sandbox, response.service_urls)
        if sandbox_ref.id == "":
            raise SandboxError("CreateSandbox returned empty sandbox id")
        return sandbox_ref

    def create_from_template(
        self,
        *,
        workspace: str,
        workload_template: str,
        spec: openshell_pb2.SandboxSpec | None = None,
        name: str | None = None,
        labels: Mapping[str, str] | None = None,
        service_exposures: Sequence[ServiceExposure] | None = None,
    ) -> SandboxRef:
        if not workload_template.strip():
            raise SandboxError("workload_template is required")
        request_spec = spec if spec is not None else openshell_pb2.SandboxSpec()
        response = self._stub.CreateSandbox(
            openshell_pb2.CreateSandboxRequest(
                workspace_scope=_workspace_scope(workspace),
                spec=request_spec,
                name=name or "",
                labels=dict(labels) if labels else {},
                workload_template=workload_template,
                service_exposures=_service_exposure_messages(service_exposures),
            ),
            timeout=self._timeout,
        )
        sandbox_ref = _sandbox_ref(response.sandbox, response.service_urls)
        if sandbox_ref.id == "":
            raise SandboxError("CreateSandbox returned empty sandbox id")
        return sandbox_ref

    def create_session(
        self,
        *,
        workspace: str,
        spec: openshell_pb2.SandboxSpec | None = None,
        name: str | None = None,
        labels: Mapping[str, str] | None = None,
        service_exposures: Sequence[ServiceExposure] | None = None,
    ) -> SandboxSession:
        return SandboxSession(
            self,
            self.create(
                workspace=workspace,
                spec=spec,
                name=name,
                labels=labels,
                service_exposures=service_exposures,
            ),
        )

    def create_session_from_template(
        self,
        *,
        workspace: str,
        workload_template: str,
        spec: openshell_pb2.SandboxSpec | None = None,
        name: str | None = None,
        labels: Mapping[str, str] | None = None,
        service_exposures: Sequence[ServiceExposure] | None = None,
    ) -> SandboxSession:
        return SandboxSession(
            self,
            self.create_from_template(
                workspace=workspace,
                workload_template=workload_template,
                spec=spec,
                name=name,
                labels=labels,
                service_exposures=service_exposures,
            ),
        )

    def sandbox_templates(self) -> SandboxTemplateClient:
        return SandboxTemplateClient(self._channel, timeout=self._timeout)

    def providers(self) -> ProviderClient:
        """Reuse this transport for provider operations; no sandbox is required."""
        from .providers import ProviderClient

        return ProviderClient(self._channel, timeout=self._timeout)

    def get(self, name: str, *, workspace: str) -> SandboxRef:
        response = self._stub.GetSandbox(
            _get_sandbox_request(name, workspace),
            timeout=self._timeout,
        )
        return _sandbox_ref(response.sandbox)

    def get_session(self, name: str, *, workspace: str) -> SandboxSession:
        return SandboxSession(self, self.get(name, workspace=workspace))

    def list(
        self,
        *,
        workspace: str,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str | None = None,
    ) -> Pager[SandboxRef]:
        def fetch(token: str) -> Page[SandboxRef]:
            response = self._stub.ListSandboxes(
                openshell_pb2.ListSandboxesRequest(
                    workspace_scope=_workspace_scope(workspace),
                    page_size=page_size,
                    page_token=token,
                    label_selector=label_selector or "",
                ),
                timeout=self._timeout,
            )
            return Page(
                items=[_sandbox_ref(item) for item in response.sandboxes],
                next_page_token=getattr(response, "next_page_token", ""),
            )

        return Pager(fetch, page_token)

    def list_all(
        self,
        *,
        workspace: str,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str | None = None,
    ) -> builtins.list[SandboxRef]:
        return self.list(
            workspace=workspace,
            page_size=page_size,
            page_token=page_token,
            label_selector=label_selector,
        ).all()

    def list_for_all_workspaces(
        self,
        *,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str | None = None,
    ) -> Pager[SandboxRef]:
        def fetch(token: str) -> Page[SandboxRef]:
            response = self._stub.ListSandboxes(
                openshell_pb2.ListSandboxesRequest(
                    workspace_scope=_all_workspaces_scope(),
                    page_size=page_size,
                    page_token=token,
                    label_selector=label_selector or "",
                ),
                timeout=self._timeout,
            )
            return Page(
                items=[_sandbox_ref(item) for item in response.sandboxes],
                next_page_token=getattr(response, "next_page_token", ""),
            )

        return Pager(fetch, page_token)

    def list_all_for_all_workspaces(
        self,
        *,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str | None = None,
    ) -> builtins.list[SandboxRef]:
        return self.list_for_all_workspaces(
            page_size=page_size,
            page_token=page_token,
            label_selector=label_selector,
        ).all()

    def list_ids(
        self,
        *,
        workspace: str,
        page_size: int = 100,
        label_selector: str | None = None,
    ) -> builtins.list[str]:
        return [
            item.id
            for item in self.list_all(
                workspace=workspace,
                page_size=page_size,
                label_selector=label_selector,
            )
        ]

    def list_ids_for_all_workspaces(
        self,
        *,
        page_size: int = 100,
        label_selector: str | None = None,
    ) -> builtins.list[str]:
        return [
            item.id
            for item in self.list_all_for_all_workspaces(
                page_size=page_size,
                label_selector=label_selector,
            )
        ]

    def delete(
        self, name: str, *, workspace: str, allow_missing: bool = False
    ) -> DeletionResult:
        response = self._stub.DeleteSandbox(
            openshell_pb2.DeleteSandboxRequest(
                workspace_scope=_workspace_scope(workspace),
                name=name,
                allow_missing=allow_missing,
            ),
            timeout=self._timeout,
        )
        return DeletionResult(
            DeletionOutcome(response.outcome), response.sandbox_id or None
        )

    def stop(self, name: str, *, workspace: str) -> SandboxRef:
        response = self._stub.StopSandbox(
            openshell_pb2.StopSandboxRequest(
                workspace_scope=_workspace_scope(workspace),
                name=name,
            ),
            timeout=self._timeout,
        )
        return _sandbox_ref(response.sandbox)

    def start(self, name: str, *, workspace: str) -> SandboxRef:
        response = self._stub.StartSandbox(
            openshell_pb2.StartSandboxRequest(
                workspace_scope=_workspace_scope(workspace),
                name=name,
            ),
            timeout=self._timeout,
        )
        return _sandbox_ref(response.sandbox)

    def suspend(self, name: str, *, workspace: str) -> SandboxRef:
        """Suspend a ready sandbox, retaining its memory state.

        Requires a compute driver and isolation backend that support suspend
        and resume.
        """
        response = self._stub.SuspendSandbox(
            openshell_pb2.SuspendSandboxRequest(
                workspace_scope=_workspace_scope(workspace),
                name=name,
            ),
            timeout=self._timeout,
        )
        return _sandbox_ref(response.sandbox)

    def resume(
        self,
        name: str,
        *,
        workspace: str,
        restart_if_restore_fails: bool = False,
    ) -> SandboxRef:
        """Resume a suspended sandbox in a new runtime generation."""
        response = self._stub.ResumeSandbox(
            openshell_pb2.ResumeSandboxRequest(
                workspace_scope=_workspace_scope(workspace),
                name=name,
                restart_if_restore_fails=restart_if_restore_fails,
            ),
            timeout=self._timeout,
        )
        return _sandbox_ref(response.sandbox)

    def wait_deleted(
        self,
        name: str,
        *,
        workspace: str,
        timeout_seconds: float = 60.0,
        expected_sandbox_id: str | None = None,
    ) -> None:
        deadline = time.time() + timeout_seconds
        while time.time() < deadline:
            try:
                current = self.get(name, workspace=workspace)
                if (
                    expected_sandbox_id is not None
                    and current.id != expected_sandbox_id
                ):
                    return
            except grpc.RpcError as exc:
                if getattr(exc, "code", lambda: None)() == grpc.StatusCode.NOT_FOUND:
                    return
                raise
            time.sleep(1)
        raise SandboxError(f"sandbox {name} was not deleted within timeout")

    def wait_ready(
        self, name: str, *, workspace: str, timeout_seconds: float = 300.0
    ) -> SandboxRef:
        return self._wait_for_phase(
            name,
            workspace=workspace,
            target_phase=openshell_pb2.SANDBOX_PHASE_READY,
            target_name="ready",
            timeout_seconds=timeout_seconds,
        )

    def wait_stopped(
        self, name: str, *, workspace: str, timeout_seconds: float = 300.0
    ) -> SandboxRef:
        return self._wait_for_phase(
            name,
            workspace=workspace,
            target_phase=openshell_pb2.SANDBOX_PHASE_STOPPED,
            target_name="stopped",
            timeout_seconds=timeout_seconds,
        )

    def wait_suspended(
        self, name: str, *, workspace: str, timeout_seconds: float = 300.0
    ) -> SandboxRef:
        return self._wait_for_phase(
            name,
            workspace=workspace,
            target_phase=openshell_pb2.SANDBOX_PHASE_SUSPENDED,
            target_name="suspended",
            timeout_seconds=timeout_seconds,
        )

    def _wait_for_phase(
        self,
        name: str,
        *,
        workspace: str,
        target_phase: int,
        target_name: str,
        timeout_seconds: float,
    ) -> SandboxRef:
        deadline = time.time() + timeout_seconds
        while time.time() < deadline:
            sandbox = self.get(name, workspace=workspace)
            if sandbox.status.phase == target_phase:
                return sandbox
            if (
                target_phase == openshell_pb2.SANDBOX_PHASE_READY
                and sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_COMPLETED
            ):
                return sandbox
            if (
                target_phase == openshell_pb2.SANDBOX_PHASE_READY
                and sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_STOPPED
            ):
                raise SandboxError(f"sandbox {name} stopped before becoming ready")
            if target_phase == openshell_pb2.SANDBOX_PHASE_READY and (
                sandbox.status.phase
                in (
                    openshell_pb2.SANDBOX_PHASE_SUSPENDING,
                    openshell_pb2.SANDBOX_PHASE_SUSPENDED,
                )
            ):
                raise SandboxError(
                    f"sandbox {name} is suspended; resume it before waiting for ready"
                )
            if sandbox.status.phase == openshell_pb2.SANDBOX_PHASE_ERROR:
                raise SandboxError(f"sandbox {name} entered error phase")
            time.sleep(1)
        raise SandboxError(f"sandbox {name} was not {target_name} within timeout")

    def exec_stream(
        self,
        sandbox: str,
        command: Sequence[str],
        *,
        workspace: 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,
    ) -> Iterator[ExecChunk | ExecResult]:
        if not command:
            raise SandboxError("command must not be empty")

        request = openshell_pb2.ExecSandboxRequest(
            workspace_scope=_workspace_scope(workspace),
            sandbox=sandbox,
            command=list(command),
            workdir=workdir or "",
            environment=dict(env or {}),
            stdin=stdin or b"",
            no_login_shell=no_login_shell,
        )
        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 = 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)

        if exit_code is None:
            raise SandboxError("ExecSandbox stream ended without an exit event")

        yield ExecResult(
            exit_code=exit_code,
            stdout=b"".join(stdout_parts).decode("utf-8", errors="replace"),
            stderr=b"".join(stderr_parts).decode("utf-8", errors="replace"),
        )

    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(
                _get_sandbox_request(sandbox, workspace),
                timeout=self._timeout,
                metadata=((_OWNER_REQUEST_HEADER, "1"),),
            )
        except grpc.RpcError:
            return ""
        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

    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,
        command: Sequence[str],
        *,
        workspace: str,
        stream_output: bool = False,
        workdir: str | None = None,
        env: Mapping[str, str] | None = None,
        stdin: bytes | None = None,
        timeout_seconds: int | None = None,
        no_login_shell: bool = False,
    ) -> ExecResult:
        result: ExecResult | None = None
        for item in self.exec_stream(
            sandbox,
            command,
            workspace=workspace,
            workdir=workdir,
            env=env,
            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
        if result is None:
            raise SandboxError("ExecSandbox did not return a result")
        return result

    def exec_python(
        self,
        sandbox: str,
        function: Callable[..., object],
        *,
        workspace: str,
        args: Sequence[object] = (),
        kwargs: Mapping[str, object] | None = None,
        stream_output: bool = False,
        workdir: str | None = None,
        env: Mapping[str, str] | None = None,
        timeout_seconds: int | None = None,
    ) -> ExecResult:
        exec_env = dict(env or {})
        exec_env["OPENSHELL_PYFUNC_B64"] = _serialize_python_callable(
            function,
            args=args,
            kwargs=kwargs,
        )
        return self.exec(
            sandbox,
            [_SANDBOX_PYTHON_BIN, "-c", _PYTHON_CLOUDPICKLE_BOOTSTRAP],
            workspace=workspace,
            stream_output=stream_output,
            workdir=workdir,
            env=exec_env,
            timeout_seconds=timeout_seconds,
        )


class SandboxTemplateClient:
    """gRPC client for reusable sandbox template lifecycle operations."""

    def __init__(self, channel: grpc.Channel, *, timeout: float = 30.0) -> None:
        self._stub = openshell_pb2_grpc.OpenShellStub(_error_mapping_channel(channel))
        self._timeout = timeout

    @classmethod
    def from_sandbox_client(cls, client: SandboxClient) -> SandboxTemplateClient:
        return client.sandbox_templates()

    def create(
        self,
        *,
        workspace: str,
        template: openshell_pb2.SandboxWorkloadTemplate | None = None,
        name: str | None = None,
        image: str | None = None,
        labels: Mapping[str, str] | None = None,
        annotations: Mapping[str, str] | None = None,
        environment: Mapping[str, str] | None = None,
        cpu: str | None = None,
        memory: str | None = None,
        gpu_count: int | None = None,
        gpu: bool = False,
        driver_config: Mapping[str, Any] | None = None,
    ) -> openshell_pb2.SandboxWorkloadTemplate:
        if template is None:
            if name is None or not name.strip():
                raise SandboxError("name is required when template is omitted")
            template = _sandbox_workload_template(
                name=name,
                image=image,
                labels=labels,
                annotations=annotations,
                environment=environment,
                cpu=cpu,
                memory=memory,
                gpu_count=gpu_count,
                gpu=gpu,
                driver_config=driver_config,
            )
        elif (
            any(value is not None for value in (name, image, cpu, memory, gpu_count))
            or any(
                bool(value)
                for value in (labels, annotations, environment, driver_config)
            )
            or gpu
        ):
            raise SandboxError(
                "template cannot be combined with template builder fields"
            )

        response = self._stub.CreateSandboxTemplate(
            openshell_pb2.CreateSandboxTemplateRequest(
                workspace_scope=_workspace_scope(workspace),
                template=template,
            ),
            timeout=self._timeout,
        )
        return response.template

    def get(
        self,
        name: str,
        *,
        workspace: str,
    ) -> openshell_pb2.SandboxWorkloadTemplate:
        response = self._stub.GetSandboxTemplate(
            openshell_pb2.GetSandboxTemplateRequest(
                workspace_scope=_workspace_scope(workspace), name=name
            ),
            timeout=self._timeout,
        )
        return response.template

    def list(
        self,
        *,
        workspace: str,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str = "",
    ) -> Pager[openshell_pb2.SandboxWorkloadTemplate]:
        def fetch(token: str) -> Page[openshell_pb2.SandboxWorkloadTemplate]:
            response = self._stub.ListSandboxTemplates(
                openshell_pb2.ListSandboxTemplatesRequest(
                    workspace_scope=_workspace_scope(workspace),
                    page_size=page_size,
                    page_token=token,
                    label_selector=label_selector,
                ),
                timeout=self._timeout,
            )
            return Page(
                items=list(response.templates),
                next_page_token=getattr(response, "next_page_token", ""),
            )

        return Pager(fetch, page_token)

    def list_all(
        self,
        *,
        workspace: str,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str = "",
    ) -> builtins.list[openshell_pb2.SandboxWorkloadTemplate]:
        return self.list(
            workspace=workspace,
            page_size=page_size,
            page_token=page_token,
            label_selector=label_selector,
        ).all()

    def list_for_all_workspaces(
        self,
        *,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str = "",
    ) -> Pager[openshell_pb2.SandboxWorkloadTemplate]:
        def fetch(token: str) -> Page[openshell_pb2.SandboxWorkloadTemplate]:
            response = self._stub.ListSandboxTemplates(
                openshell_pb2.ListSandboxTemplatesRequest(
                    workspace_scope=_all_workspaces_scope(),
                    page_size=page_size,
                    page_token=token,
                    label_selector=label_selector,
                ),
                timeout=self._timeout,
            )
            return Page(
                items=list(response.templates),
                next_page_token=getattr(response, "next_page_token", ""),
            )

        return Pager(fetch, page_token)

    def list_all_for_all_workspaces(
        self,
        *,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str = "",
    ) -> builtins.list[openshell_pb2.SandboxWorkloadTemplate]:
        return self.list_for_all_workspaces(
            page_size=page_size,
            page_token=page_token,
            label_selector=label_selector,
        ).all()

    def delete(
        self, name: str, *, workspace: str, allow_missing: bool = False
    ) -> DeletionResult:
        response = self._stub.DeleteSandboxTemplate(
            openshell_pb2.DeleteSandboxTemplateRequest(
                workspace_scope=_workspace_scope(workspace),
                name=name,
                allow_missing=allow_missing,
            ),
            timeout=self._timeout,
        )
        return DeletionResult(DeletionOutcome(response.outcome))


@dataclass(frozen=True)
class WorkspaceRef:
    name: str
    phase: str
    labels: dict[str, str]


def _workspace_ref(ws: datamodel_pb2.Workspace) -> WorkspaceRef:
    meta = ws.metadata
    return WorkspaceRef(
        name=meta.name,
        phase=datamodel_pb2.WorkspacePhase.Name(ws.status.phase),
        labels=dict(meta.labels),
    )


class WorkspaceClient:
    """gRPC client for workspace lifecycle operations."""

    def __init__(self, channel: grpc.Channel, *, timeout: float = 30.0) -> None:
        self._stub = openshell_pb2_grpc.OpenShellStub(_error_mapping_channel(channel))
        self._timeout = timeout

    @classmethod
    def from_sandbox_client(cls, client: SandboxClient) -> WorkspaceClient:
        return cls(client._channel, timeout=client._timeout)

    def create(
        self,
        name: str,
        *,
        labels: Mapping[str, str] | None = None,
    ) -> WorkspaceRef:
        response = self._stub.CreateWorkspace(
            openshell_pb2.CreateWorkspaceRequest(
                name=name,
                labels=dict(labels) if labels else {},
            ),
            timeout=self._timeout,
        )
        return _workspace_ref(response.workspace)

    def get(self, name: str) -> WorkspaceRef:
        response = self._stub.GetWorkspace(
            openshell_pb2.GetWorkspaceRequest(name=name),
            timeout=self._timeout,
        )
        return _workspace_ref(response.workspace)

    def list(
        self,
        *,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str | None = None,
    ) -> Pager[WorkspaceRef]:
        def fetch(token: str) -> Page[WorkspaceRef]:
            response = self._stub.ListWorkspaces(
                openshell_pb2.ListWorkspacesRequest(
                    page_size=page_size,
                    page_token=token,
                    label_selector=label_selector or "",
                ),
                timeout=self._timeout,
            )
            return Page(
                items=[_workspace_ref(ws) for ws in response.workspaces],
                next_page_token=getattr(response, "next_page_token", ""),
            )

        return Pager(fetch, page_token)

    def list_all(
        self,
        *,
        page_size: int = 100,
        page_token: str = "",
        label_selector: str | None = None,
    ) -> builtins.list[WorkspaceRef]:
        return self.list(
            page_size=page_size,
            page_token=page_token,
            label_selector=label_selector,
        ).all()

    def delete(self, name: str, *, allow_missing: bool = False) -> DeletionResult:
        response = self._stub.DeleteWorkspace(
            openshell_pb2.DeleteWorkspaceRequest(
                name=name, allow_missing=allow_missing
            ),
            timeout=self._timeout,
        )
        return DeletionResult(DeletionOutcome(response.outcome))


class Sandbox:
    """Context-managed sandbox session bound to one sandbox id."""

    def __init__(
        self,
        *,
        workspace: str,
        cluster: str | None = None,
        sandbox: str | SandboxRef | None = None,
        delete_on_exit: bool = True,
        spec: openshell_pb2.SandboxSpec | None = None,
        name: str | None = None,
        labels: Mapping[str, str] | None = None,
        workload_template: str | None = None,
        timeout: float = 30.0,
        ready_timeout_seconds: float = 120.0,
        auto_refresh: bool = True,
        write_back: bool = True,
        insecure: bool = False,
        client_credentials: ClientCredentialsAuth | None = None,
    ) -> None:
        """Bind a Sandbox context to the active gateway.

        OIDC kwargs (`auto_refresh`, `write_back`, `insecure`, and
        `client_credentials`) forward
        directly to `SandboxClient.from_active_cluster` and have the
        same semantics. They're surfaced on `Sandbox` so callers using
        the higher-level wrapper get parity with `SandboxClient` for
        OIDC-protected gateways (e.g. passing `insecure=True` for a
        self-signed dev IdP). Non-OIDC gateways ignore them.
        """
        self._workspace = workspace
        self._cluster = cluster
        self._sandbox_input = sandbox
        self._delete_on_exit = delete_on_exit
        self._spec = spec
        self._name = name
        # Copy so later caller mutation cannot change what gets sent on enter.
        self._labels = dict(labels) if labels is not None else None
        self._workload_template = workload_template
        self._timeout = timeout
        self._ready_timeout_seconds = ready_timeout_seconds
        self._auto_refresh = auto_refresh
        self._write_back = write_back
        self._insecure = insecure
        self._client_credentials = client_credentials
        self._client: SandboxClient | None = None
        self._session: SandboxSession | None = None

    @property
    def id(self) -> str:
        if self._session is None:
            raise SandboxError("sandbox context has not been entered")
        return self._session.id

    @property
    def sandbox(self) -> SandboxRef:
        if self._session is None:
            raise SandboxError("sandbox context has not been entered")
        return self._session.sandbox

    def __enter__(self) -> Sandbox:
        # Creation metadata cannot be applied when attaching to an existing
        # sandbox; reject it before opening a connection.
        if self._sandbox_input is not None and (
            self._name is not None
            or self._labels is not None
            or self._workload_template is not None
        ):
            raise SandboxError(
                "name, labels, and workload_template cannot be set when attaching to an existing sandbox"
            )

        self._client = SandboxClient.from_active_cluster(
            cluster=self._cluster,
            timeout=self._timeout,
            auto_refresh=self._auto_refresh,
            write_back=self._write_back,
            insecure=self._insecure,
            client_credentials=self._client_credentials,
        )
        try:
            self._open_session(self._client)
        except BaseException:
            # Python skips __exit__ when __enter__ raises, so release here
            # and keep the original error even if the close fails.
            with contextlib.suppress(Exception):
                self._release()
            raise
        return self

    def _open_session(self, client: SandboxClient) -> None:
        if self._sandbox_input is None and self._workload_template is not None:
            self._session = client.create_session_from_template(
                workspace=self._workspace,
                workload_template=self._workload_template,
                spec=self._spec,
                name=self._name,
                labels=self._labels,
            )
        elif self._sandbox_input is None:
            self._session = client.create_session(
                workspace=self._workspace,
                spec=self._spec,
                name=self._name,
                labels=self._labels,
            )
        elif isinstance(self._sandbox_input, SandboxRef):
            self._session = SandboxSession(client, self._sandbox_input)
        else:
            self._session = client.get_session(
                self._sandbox_input, workspace=self._workspace
            )

        self._workspace = getattr(self._session, "_workspace", self._workspace)

        ready = client.wait_ready(
            self._session.sandbox.name,
            workspace=self._workspace,
            timeout_seconds=self._ready_timeout_seconds,
        )
        self._session = SandboxSession(client, ready)

    def _release(self) -> None:
        client = self._client
        self._session = None
        self._client = None
        if client is not None:
            client.close()

    def __exit__(self, *args: object) -> None:
        try:
            if (
                self._delete_on_exit
                and self._session is not None
                and self._client is not None
            ):
                result = self._session.delete(allow_missing=True)
                if result.outcome == DeletionOutcome.ACCEPTED:
                    self._client.wait_deleted(
                        self._session.sandbox.name,
                        workspace=self._workspace,
                        expected_sandbox_id=result.sandbox_id,
                    )
                elif result.outcome not in (
                    DeletionOutcome.COMPLETED,
                    DeletionOutcome.ALREADY_ABSENT,
                ):
                    raise SandboxError(
                        f"unsupported deletion outcome: {result.outcome}"
                    )
        finally:
            self._release()

    def exec(
        self,
        command: Sequence[str],
        *,
        stream_output: bool = False,
        workdir: str | None = None,
        env: Mapping[str, str] | None = None,
        stdin: bytes | None = None,
        timeout_seconds: int | None = None,
        no_login_shell: bool = False,
    ) -> ExecResult:
        if self._session is None:
            raise SandboxError("sandbox context has not been entered")
        return self._session.exec(
            command,
            stream_output=stream_output,
            workdir=workdir,
            env=env,
            stdin=stdin,
            timeout_seconds=timeout_seconds,
            no_login_shell=no_login_shell,
        )

    def exec_python(
        self,
        function: Callable[..., object],
        *,
        args: Sequence[object] = (),
        kwargs: Mapping[str, object] | None = None,
        stream_output: bool = False,
        workdir: str | None = None,
        env: Mapping[str, str] | None = None,
        timeout_seconds: int | None = None,
    ) -> ExecResult:
        if self._session is None:
            raise SandboxError("sandbox context has not been entered")
        return self._session.exec_python(
            function,
            args=args,
            kwargs=kwargs,
            stream_output=stream_output,
            workdir=workdir,
            env=env,
            timeout_seconds=timeout_seconds,
        )


_PYTHON_CLOUDPICKLE_BOOTSTRAP = (
    "import base64,cloudpickle,os;"
    "payload=base64.b64decode(os.environ['OPENSHELL_PYFUNC_B64']);"
    "func,args,kwargs=cloudpickle.loads(payload);"
    "result=func(*args,**kwargs);"
    "print(result) if result is not None else None"
)

_SANDBOX_PYTHON_BIN = "python"


def _serialize_python_callable(
    function: Callable[..., object],
    *,
    args: Sequence[object],
    kwargs: Mapping[str, object] | None,
) -> str:
    try:
        import cloudpickle
    except ImportError as exc:  # pragma: no cover - import error path
        raise SandboxError("cloudpickle is required for exec_python") from exc

    payload = cloudpickle.dumps((function, tuple(args), dict(kwargs or {})))
    return base64.b64encode(payload).decode("ascii")


def _sandbox_conditions(
    status: openshell_pb2.SandboxStatus | None,
) -> tuple[SandboxConditionRef, ...]:
    if status is None:
        return ()
    return tuple(
        SandboxConditionRef(
            type=condition.type,
            status=condition.status,
            reason=condition.reason,
            message=condition.message,
            transition_time_ms=condition.transition_time.ToMilliseconds()
            if condition.HasField("transition_time")
            else None,
        )
        for condition in status.conditions
    )


def _sandbox_ref(
    sandbox: openshell_pb2.Sandbox,
    service_urls: Mapping[str, str] | None = None,
) -> SandboxRef:
    status = sandbox.status if sandbox.HasField("status") else None
    provenance = (
        SandboxWorkloadTemplateProvenanceRef(
            name=sandbox.created_from_workload_template.name,
            resource_version=sandbox.created_from_workload_template.resource_version,
        )
        if sandbox.HasField("created_from_workload_template")
        else None
    )
    return SandboxRef(
        id=sandbox.metadata.id if sandbox.metadata else "",
        name=sandbox.metadata.name if sandbox.metadata else "",
        workspace=sandbox.metadata.workspace if sandbox.metadata else "",
        status=SandboxStatusRef(
            phase=status.phase if status else 0,
            current_policy_version=status.current_policy_version if status else 0,
            exit_code=status.exit_code
            if status is not None and status.HasField("exit_code")
            else None,
            restart_count=status.restart_count if status is not None else 0,
            next_restart_at_ms=status.next_restart_time.ToMilliseconds()
            if status is not None and status.HasField("next_restart_time")
            else None,
            main_process_started_at_ms=status.main_process_started_time.ToMilliseconds()
            if status is not None and status.HasField("main_process_started_time")
            else None,
            conditions=_sandbox_conditions(status),
        ),
        labels=sandbox.metadata.labels if sandbox.metadata else {},
        created_from_workload_template=provenance,
        service_urls=service_urls or {},
        host_key_fingerprint=sandbox.host_key_fingerprint or None,
    )


def _default_spec() -> openshell_pb2.SandboxSpec:
    # Omit the policy field so the sandbox container discovers its policy
    # from /etc/openshell/policy.yaml (baked into the image at build time).
    # This avoids duplicating policy defaults between the SDK and the
    # container image and ensures sandboxes get the full dev-sandbox-policy
    # (including network_policies) out of the box.
    return openshell_pb2.SandboxSpec()


def _sandbox_workload_template(
    *,
    name: str,
    image: str | None = None,
    labels: Mapping[str, str] | None = None,
    annotations: Mapping[str, str] | None = None,
    environment: Mapping[str, str] | None = None,
    cpu: str | None = None,
    memory: str | None = None,
    gpu_count: int | None = None,
    gpu: bool = False,
    driver_config: Mapping[str, Any] | None = None,
) -> openshell_pb2.SandboxWorkloadTemplate:
    if gpu_count is not None and gpu_count <= 0:
        raise SandboxError("gpu_count must be greater than zero")

    template = openshell_pb2.SandboxWorkloadTemplate()
    template.metadata.name = name
    # The gateway permits an empty workload so it can apply the default image
    # at sandbox create time, but the workload message itself is required.
    template.spec.workload.SetInParent()
    if labels:
        template.metadata.labels.update(dict(labels))
    if annotations:
        template.metadata.annotations.update(dict(annotations))

    if image is not None:
        template.spec.workload.image = image
    if environment:
        template.spec.workload.environment.update(dict(environment))
    if cpu is not None:
        template.spec.workload.resources.cpu = cpu
    if memory is not None:
        template.spec.workload.resources.memory = memory
    if gpu or gpu_count is not None:
        template.spec.workload.resources.gpu.SetInParent()
    if gpu_count is not None:
        template.spec.workload.resources.gpu.count = gpu_count
    if driver_config:
        template.spec.driver_config.update(dict(driver_config))

    return template


def _xdg_config_home() -> pathlib.Path:
    configured = os.environ.get("XDG_CONFIG_HOME")
    if configured:
        return pathlib.Path(configured)
    return pathlib.Path.home() / ".config"


# Re-check the cached token roughly 30 seconds before the issuer's
# stated expiry, to leave room for in-flight RPCs and clock skew. This
# matches `openshell-bootstrap::oidc_token::is_token_expired`.
_OIDC_TOKEN_EXPIRY_GRACE_SECONDS = 30

_IS_WINDOWS = os.name == "nt"
_WINDOWS_REPLACE_RETRYABLE_ERRORS = frozenset({5, 32})
_WINDOWS_REPLACE_TIMEOUT_SECONDS = 0.25
_WINDOWS_REPLACE_INITIAL_DELAY_SECONDS = 0.005
_WINDOWS_REPLACE_MAX_DELAY_SECONDS = 0.05
_WINDOWS_REPLACE_LOCK = threading.Lock()


def _atomic_replace(source: pathlib.Path, destination: pathlib.Path) -> None:
    """Atomically replace a file, retrying transient Windows sharing errors."""
    if not _IS_WINDOWS:
        source.replace(destination)
        return

    # Serialize writers in this process. The retry still handles other
    # processes (including the Rust CLI) and filesystem scanners that briefly
    # open the destination without delete sharing.
    with _WINDOWS_REPLACE_LOCK:
        deadline = time.monotonic() + _WINDOWS_REPLACE_TIMEOUT_SECONDS
        delay = _WINDOWS_REPLACE_INITIAL_DELAY_SECONDS
        while True:
            try:
                source.replace(destination)
                return
            except PermissionError as error:
                winerror = getattr(error, "winerror", None)
                retryable = winerror in _WINDOWS_REPLACE_RETRYABLE_ERRORS or (
                    winerror is None and error.errno == errno.EACCES
                )
                if not retryable or time.monotonic() >= deadline:
                    raise
                time.sleep(delay)
                delay = min(delay * 2, _WINDOWS_REPLACE_MAX_DELAY_SECONDS)


def _read_oidc_token_bundle(gateway_dir: pathlib.Path) -> dict | None:
    """Read and parse `oidc_token.json` for a gateway.

    Returns the parsed dict, or `None` if the file is absent or unreadable.
    See `openshell-bootstrap::oidc_token::store_oidc_token` for the writer.
    """
    token_path = gateway_dir / "oidc_token.json"
    try:
        return json.loads(token_path.read_text(encoding="utf-8"))
    except FileNotFoundError:
        return None
    except (OSError, UnicodeDecodeError, json.JSONDecodeError):
        return None


def _normalize_issuer(bundle: dict) -> str | None:
    """Return the bundle's issuer with a trailing slash stripped.

    Used to detect whether the issuer changed when adopting a bundle
    re-read from disk, so a cached token endpoint computed for the old
    issuer can be invalidated. Trailing-slash differences are treated as
    equal, matching `_discover_token_endpoint`'s normalization.
    """
    issuer = bundle.get("issuer")
    return issuer.rstrip("/") if isinstance(issuer, str) else None


def _load_cluster_bearer_token(gateway_dir: pathlib.Path) -> str | None:
    """Read a single (possibly expired) access token from disk.

    Lower-level helper used by both the legacy single-shot path and the
    refreshing provider. Returns the raw access_token string or None.
    """
    bundle = _read_oidc_token_bundle(gateway_dir)
    if bundle is None:
        return None
    access_token = bundle.get("access_token")
    if not isinstance(access_token, str) or not access_token:
        return None
    return access_token


def _make_fail_closed_bearer_provider(
    gateway_dir: pathlib.Path,
    cluster_name: str,
) -> Callable[[], str]:
    """Per-RPC provider that re-reads `oidc_token.json` but does NOT refresh.

    Raises `SandboxError` when the token is missing or expired. Available as
    an opt-out from the default `_OidcRefresher` for callers (e.g. tests)
    that want to assert expiry behavior or that don't want the SDK to make
    outbound HTTP calls to the IdP.
    """

    def provider() -> str:
        bundle = _read_oidc_token_bundle(gateway_dir)
        if bundle is None:
            raise SandboxError(
                f"OIDC token for gateway '{cluster_name}' is missing or "
                f"unreadable. Re-authenticate with: openshell gateway login"
            )
        access_token = bundle.get("access_token")
        if not isinstance(access_token, str) or not access_token:
            raise SandboxError(
                f"OIDC token for gateway '{cluster_name}' has no access "
                f"token. Re-authenticate with: openshell gateway login"
            )
        expires_at = bundle.get("expires_at")
        if isinstance(expires_at, int):
            now = int(time.time())
            if now + _OIDC_TOKEN_EXPIRY_GRACE_SECONDS >= expires_at:
                raise SandboxError(
                    f"OIDC token for gateway '{cluster_name}' has expired. "
                    f"Re-authenticate with: openshell gateway login"
                )
        return access_token

    return provider


class _InvalidGrantError(SandboxError):
    """Refresh failed with OAuth2 `invalid_grant` (RFC 6749 §5.2).

    The refresh_token was rejected — expired, revoked, or rotated out by
    a concurrent refresh in another process. Subclass of `SandboxError`
    so an uncaught instance still surfaces the standard re-authenticate
    hint; `current_access_token` catches it to retry once with a
    peer-rotated bundle before giving up.
    """


class _OidcRefresher:
    """Thread-safe in-process OAuth2 refresh for a gateway's `oidc_token.json`.

    Mirrors the lazy-refresh pattern used by `google-auth`'s `Credentials`
    and `botocore`'s `SSOTokenProvider`. Uses `httpx` for transport so
    we can pin `follow_redirects=False` and the TLS verification policy
    explicitly — same posture as the Rust CLI's use of `reqwest` with
    `Policy::none()` and an opt-in `danger_accept_invalid_certs` (see
    `crates/openshell-cli/src/oidc_auth.rs::http_client`). The OAuth2
    refresh-token grant itself (RFC 6749 §6) is a single form-encoded
    POST, handled inline rather than via an OAuth2 library.

    Properties:

    - **Lazy**: check expiry on every RPC; refresh only when stale.
    - **Lock-coordinated**: concurrent RPCs share a single refresh, not
      one per call. Plain `threading.Lock` (no separate worker thread,
      unlike `google-auth`'s `RefreshThreadManager` — sufficient for
      our use case).
    - **Disk-aware**: before refreshing, re-read `oidc_token.json` —
      the CLI or another process may have already rotated the bundle.
    - **Discovery-validated**: fetches the OIDC discovery document
      from `<issuer>/.well-known/openid-configuration`, rejects
      responses whose `issuer` field doesn't match the configured one
      (preventing SSRF / misdirection of the refresh_token to an
      attacker-controlled endpoint). Mirrors the Rust CLI's
      `discover()` function.
    - **Redirect-hardened**: `follow_redirects=False` on the underlying
      `httpx.Client` so a 3xx during discovery or refresh is treated
      as a failure rather than silently chasing the redirect to an
      arbitrary host.
    - **Write-back by default**: when `auto_refresh=True`, refreshed
      bundles are atomically persisted to `oidc_token.json` at mode
      0600 so other processes (Rust CLI, TUI, other Python clients)
      see the rotated `refresh_token`. Required for IdPs that
      invalidate the old `refresh_token` on rotation (Keycloak with
      rotation enabled, Entra in strict mode); without write-back, a
      second process would `invalid_grant` on next refresh.
    - **`insecure=True` flag** disables TLS certificate verification
      for the discovery and refresh calls. Matches the Rust CLI's
      `--insecure` plumbing for OIDC issuers behind self-signed certs.
    - **Refresh failures** surface as `SandboxError` with a
      "re-authenticate with: openshell gateway login" hint.
    """

    def __init__(
        self,
        gateway_dir: pathlib.Path,
        cluster_name: str,
        *,
        write_back: bool = True,
        insecure: bool = False,
    ) -> None:
        self._gateway_dir = gateway_dir
        self._cluster_name = cluster_name
        self._write_back = write_back
        self._lock = threading.Lock()
        self._bundle: dict | None = None
        self._token_endpoint: str | None = None
        # Single httpx.Client serves both discovery (unauthenticated
        # GET) and refresh (POST with form-encoded body) — they share
        # the same security posture.
        #
        # - follow_redirects=False: a 3xx during discovery would
        #   otherwise steer us to an attacker-controlled token
        #   endpoint. Matches `reqwest::redirect::Policy::none()` in
        #   the Rust CLI's `oidc_auth.rs::http_client`.
        # - verify=not insecure: opt-in TLS-verification disable for
        #   self-signed issuers. Matches the Rust CLI's `--insecure`
        #   flag plumbing.
        #
        # We don't use authlib's `OAuth2Client` here because it
        # auto-injects an Authorization header on every request from
        # its stored token, which would break the unauthenticated
        # discovery GET. The refresh_token grant for a public client
        # is a single form-encoded POST — small enough to spell out
        # directly with httpx, and easier to test.
        self._http = httpx.Client(
            follow_redirects=False,
            verify=not insecure,
            timeout=15.0,
        )

    def close(self) -> None:
        self._http.close()

    def __del__(self) -> None:
        # Best-effort cleanup; tests + short-lived callers may not call
        # close() explicitly.
        with contextlib.suppress(Exception):
            self._http.close()

    def current_access_token(self) -> str:
        """Return a non-expired access token, refreshing if needed."""
        with self._lock:
            if self._bundle is None:
                self._bundle = _read_oidc_token_bundle(self._gateway_dir)
                if self._bundle is None:
                    raise SandboxError(
                        f"OIDC token for gateway '{self._cluster_name}' is "
                        f"missing or unreadable. Re-authenticate with: "
                        f"openshell gateway login"
                    )
            if self._is_fresh(self._bundle):
                return self._bundle["access_token"]
            # Cached bundle is stale. Before refreshing, re-read disk —
            # another process (CLI, TUI, another SDK client) may have
            # rotated the bundle while we were idle. Adopt the disk
            # bundle when it was refreshed more recently than ours, EVEN
            # WHEN its access token is also stale: otherwise we'd refresh
            # with our in-memory refresh_token, which a rotating IdP may
            # have already invalidated when the other process refreshed
            # (Keycloak with rotation, Entra in strict mode).
            #
            # "More recently" is judged by `expires_at`: a refresh issues
            # a new access token with a forward expiry alongside the
            # (possibly rotated) refresh_token, so the bundle with the
            # later expiry carries the newest refresh_token. This also
            # preserves the write_back=False case, where our in-memory
            # bundle has already rotated past the on-disk one and must
            # NOT be clobbered by the older disk copy.
            disk = _read_oidc_token_bundle(self._gateway_dir)
            if disk is not None and self._expiry(disk) > self._expiry(self._bundle):
                # If the issuer changed under us, the cached token
                # endpoint no longer applies — force re-discovery.
                if _normalize_issuer(disk) != _normalize_issuer(self._bundle):
                    self._token_endpoint = None
                self._bundle = disk
                if self._is_fresh(disk):
                    return disk["access_token"]
            # Truly stale; refresh against the IdP using the freshest
            # bundle we have (disk if it was newer, else in-memory).
            try:
                self._bundle = self._refresh(self._bundle)
            except _InvalidGrantError as exc:
                # We lost a cross-process rotation race: between our disk
                # re-read above and our refresh POST, a peer (CLI, TUI,
                # another SDK client) rotated the refresh_token and the
                # IdP invalidated ours. This is the residual window that
                # neither google-auth nor botocore close without an OS
                # file lock. Rather than lock, recover: re-read disk once
                # and, if a peer wrote a *different* refresh_token, retry
                # with it before surfacing a re-authenticate error.
                self._bundle = self._recover_from_invalid_grant(self._bundle, exc)
            if self._write_back:
                self._write_to_disk(self._bundle)
            return self._bundle["access_token"]

    def _recover_from_invalid_grant(
        self, attempted: dict, exc: _InvalidGrantError
    ) -> dict:
        """Re-read disk after an `invalid_grant` and retry once if a peer
        rotated the refresh_token.

        `attempted` is the bundle whose refresh_token the IdP just
        rejected. If disk now holds a different refresh_token, a
        concurrent process won the rotation race — adopt and reuse it
        (returning early if it is already fresh, otherwise refreshing
        with it). If disk offers nothing new, the rejection is genuine:
        re-raise so the caller sees the re-authenticate hint.
        """
        disk = _read_oidc_token_bundle(self._gateway_dir)
        if disk is None or disk.get("refresh_token") == attempted.get("refresh_token"):
            # No peer rotation — the refresh_token really is dead.
            raise exc
        if _normalize_issuer(disk) != _normalize_issuer(attempted):
            self._token_endpoint = None
        if self._is_fresh(disk):
            return disk
        # Single retry with the peer's rotated token; a second
        # _InvalidGrantError here propagates (no further retry).
        return self._refresh(disk)

    @staticmethod
    def _is_fresh(bundle: dict) -> bool:
        access_token = bundle.get("access_token")
        if not isinstance(access_token, str) or not access_token:
            return False
        exp = bundle.get("expires_at")
        if not isinstance(exp, int):
            # No expiry info — treat as fresh (matches the
            # `is_token_expired` semantics in the Rust CLI).
            return True
        return int(time.time()) + _OIDC_TOKEN_EXPIRY_GRACE_SECONDS < exp

    @staticmethod
    def _expiry(bundle: dict) -> float:
        """Access-token expiry as a comparable number.

        A bundle without an `expires_at` is treated as non-expiring
        (`+inf`) — consistent with `_is_fresh`, which treats a missing
        expiry as always fresh. Used to decide which of two stale
        bundles (in-memory vs. on-disk) was refreshed more recently and
        therefore holds the newest refresh_token.
        """
        exp = bundle.get("expires_at")
        return float(exp) if isinstance(exp, int) else float("inf")

    def _discover_token_endpoint(self, bundle: dict) -> str:
        if self._token_endpoint is not None:
            return self._token_endpoint
        issuer = bundle.get("issuer")
        if not isinstance(issuer, str) or not issuer:
            raise SandboxError(
                f"OIDC bundle for gateway '{self._cluster_name}' has no "
                f"`issuer`; cannot refresh. Re-authenticate with: openshell "
                f"gateway login"
            )
        normalized_issuer = issuer.rstrip("/")
        discovery_url = f"{normalized_issuer}/.well-known/openid-configuration"
        try:
            resp = self._http.get(discovery_url)
        except httpx.HTTPError as e:
            raise SandboxError(
                f"OIDC discovery failed for gateway "
                f"'{self._cluster_name}': {e}. Re-authenticate with: "
                f"openshell gateway login"
            ) from e
        # follow_redirects=False means a 3xx surfaces as a non-2xx
        # status; treat any non-success as a discovery failure rather
        # than silently following.
        if not 200 <= resp.status_code < 300:
            raise SandboxError(
                f"OIDC discovery failed for gateway "
                f"'{self._cluster_name}': HTTP {resp.status_code} "
                f"from {discovery_url}. Re-authenticate with: openshell "
                f"gateway login"
            )
        try:
            disco = resp.json()
        except ValueError as e:
            raise SandboxError(
                f"OIDC discovery returned invalid JSON for gateway "
                f"'{self._cluster_name}': {e}"
            ) from e
        # Critical: validate that the discovery document's `issuer`
        # matches the configured one. Without this, a misdirected or
        # malicious discovery response could steer the refresh_token
        # POST to an attacker-controlled endpoint.
        discovered_issuer = disco.get("issuer", "")
        if not isinstance(discovered_issuer, str) or (
            discovered_issuer.rstrip("/") != normalized_issuer
        ):
            raise SandboxError(
                f"OIDC discovery issuer mismatch for gateway "
                f"'{self._cluster_name}': expected '{normalized_issuer}', "
                f"got '{discovered_issuer}'."
            )
        endpoint = disco.get("token_endpoint")
        if not isinstance(endpoint, str) or not endpoint:
            raise SandboxError(
                f"OIDC discovery for gateway '{self._cluster_name}' did "
                f"not include a token_endpoint."
            )
        self._token_endpoint = endpoint
        return endpoint

    def _refresh(self, bundle: dict) -> dict:
        refresh_token = bundle.get("refresh_token")
        if not isinstance(refresh_token, str) or not refresh_token:
            raise SandboxError(
                f"OIDC token for gateway '{self._cluster_name}' has no "
                f"refresh token. Re-authenticate with: openshell gateway "
                f"login"
            )
        token_endpoint = self._discover_token_endpoint(bundle)
        client_id = bundle.get("client_id", "openshell-cli")

        # RFC 6749 §6: refresh_token grant. Form-encoded POST with
        # grant_type, refresh_token, and (for a public client) client_id.
        # No Authorization header (token_endpoint_auth_method="none").
        try:
            resp = self._http.post(
                token_endpoint,
                data={
                    "grant_type": "refresh_token",
                    "refresh_token": refresh_token,
                    "client_id": client_id,
                },
            )
        except httpx.HTTPError as e:
            raise SandboxError(
                f"OIDC token refresh failed for gateway "
                f"'{self._cluster_name}': {type(e).__name__}: {e}. "
                f"Re-authenticate with: openshell gateway login"
            ) from e
        if resp.status_code != 200:
            # Include the IdP's error body for diagnostics — RFC 6749
            # mandates a JSON body like {"error":"invalid_grant", ...}
            # on failure, which is the most useful signal to surface.
            error_code = None
            with contextlib.suppress(Exception):
                body = resp.json()
                if isinstance(body, dict):
                    error_code = body.get("error")
            detail = ""
            with contextlib.suppress(Exception):
                detail = f": {resp.text[:200]}"
            message = (
                f"OIDC token refresh failed for gateway "
                f"'{self._cluster_name}': HTTP {resp.status_code}"
                f"{detail}. Re-authenticate with: openshell gateway "
                f"login"
            )
            # `invalid_grant` specifically means the refresh_token was
            # rejected — distinguished from transport/5xx errors so the
            # caller can retry once with a peer-rotated bundle (a lost
            # cross-process rotation race) before surfacing the failure.
            if error_code == "invalid_grant":
                raise _InvalidGrantError(message)
            raise SandboxError(message)
        try:
            token = resp.json()
        except ValueError as e:
            raise SandboxError(
                f"OIDC refresh response for gateway '{self._cluster_name}' "
                f"is not JSON: {e}"
            ) from e

        access_token = token.get("access_token")
        if not isinstance(access_token, str) or not access_token:
            raise SandboxError(
                f"OIDC refresh response for gateway '{self._cluster_name}' "
                f"is missing access_token."
            )
        expires_at = token.get("expires_at")
        if expires_at is None:
            expires_in = token.get("expires_in")
            if isinstance(expires_in, (int, float)):
                expires_at = int(time.time()) + int(expires_in)
        return {
            "access_token": access_token,
            # Refresh-token rotation: some IdPs (Keycloak with rotation
            # enabled, Entra in strict mode) reissue and invalidate the
            # old one. Honor the new value when present.
            "refresh_token": token.get("refresh_token", refresh_token),
            "expires_at": int(expires_at) if expires_at is not None else None,
            "issuer": bundle.get("issuer", ""),
            "client_id": client_id,
        }

    def _write_to_disk(self, bundle: dict) -> None:
        """Atomic-replace `oidc_token.json` with the refreshed bundle.

        Strips `None` values to match the Rust writer's
        `skip_serializing_if = "Option::is_none"` behavior so a Python-
        written file is byte-identical in shape to what the CLI writes.

        Uses `tempfile.mkstemp` (PID + random suffix) so two writers
        racing on the same gateway directory don't share a tmp file
        and trample each other's content. Each writer gets a unique
        path; `.replace()` is atomic per-writer, and POSIX rename
        semantics ensure the final `oidc_token.json` is always
        complete-and-readable to anyone observing.
        """
        path = self._gateway_dir / "oidc_token.json"
        serializable = {k: v for k, v in bundle.items() if v is not None}
        payload = json.dumps(serializable, indent=2)

        # mkstemp creates the file with mode 0600 already on POSIX
        # (it uses O_CREAT | O_EXCL with restrictive umask), so chmod
        # is a belt-and-braces step for filesystems that don't honor
        # the initial mode.
        fd, tmp_name = tempfile.mkstemp(
            prefix=".oidc_token.",
            suffix=".tmp",
            dir=str(self._gateway_dir),
        )
        tmp_path = pathlib.Path(tmp_name)
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as f:
                f.write(payload)
            with contextlib.suppress(OSError):
                tmp_path.chmod(0o600)
            _atomic_replace(tmp_path, path)
        except BaseException:
            # Clean up our tmp on failure so we don't leave orphaned
            # `.oidc_token.<rand>.tmp` files lying around. The replace
            # already moved the file on the success path.
            with contextlib.suppress(OSError):
                tmp_path.unlink()
            raise


def _make_cluster_bearer_provider(
    gateway_dir: pathlib.Path,
    cluster_name: str,
    *,
    auto_refresh: bool = True,
    write_back: bool = True,
    insecure: bool = False,
) -> tuple[Callable[[], str], Callable[[], None] | None]:
    """Build a per-RPC token provider for a gateway directory.

    Returns `(token_provider, close_fn_or_none)`. `close_fn` is non-None
    only when an `_OidcRefresher` was constructed; callers that own the
    provider's lifecycle (e.g. `SandboxClient.close()`) should invoke
    it during teardown so the underlying httpx.Client is released
    rather than relying on `__del__`.

    With `auto_refresh=True` (the default), returns an `_OidcRefresher`-
    backed callable that lazily refreshes against the IdP's token endpoint
    when the cached bundle is stale. This mirrors the lazy-refresh pattern
    used by `google.oauth2.credentials.Credentials` and
    `botocore.tokens.SSOTokenProvider` and lets long-running scripts
    survive token rotation without intervention.

    With `auto_refresh=False`, falls back to the read-only / fail-closed
    behavior: the SDK consumes whatever the CLI most recently wrote and
    raises `SandboxError` when the token expires. Useful for tests or
    callers that don't want the SDK to make outbound HTTP calls to the
    IdP. No close_fn is returned in this case.

    `write_back=True` (only meaningful when `auto_refresh=True`) makes the
    refresher atomically persist the rotated bundle back to
    `oidc_token.json` so other processes — including the Rust CLI — see
    the new token. Defaults to True because OIDC providers with
    refresh-token rotation (Keycloak, Entra) invalidate the old
    refresh_token on rotation; an in-memory-only refresh would leave the
    on-disk bundle pointing at an invalidated value, and any other
    process starting from that disk state would fail on its first
    refresh.

    `insecure=True` disables TLS certificate verification for both the
    OIDC discovery document fetch and the refresh-token POST. Mirrors
    the Rust CLI's `--insecure` flag for OIDC issuers behind self-signed
    certs.
    """
    if not auto_refresh:
        return _make_fail_closed_bearer_provider(gateway_dir, cluster_name), None
    refresher = _OidcRefresher(
        gateway_dir,
        cluster_name,
        write_back=write_back,
        insecure=insecure,
    )
    return refresher.current_access_token, refresher.close


def _resolve_active_cluster() -> str:
    env_gateway = os.environ.get("OPENSHELL_GATEWAY")
    if env_gateway:
        return env_gateway
    active_file = _xdg_config_home() / "openshell" / "active_gateway"
    try:
        value = active_file.read_text(encoding="utf-8").strip()
    except FileNotFoundError:
        raise SandboxError("no active gateway configured") from None
    if value == "":
        raise SandboxError("no active gateway configured")
    return value