diff --git a/Cargo.lock b/Cargo.lock index 21c81cc945..1b93758bdc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3931,6 +3931,7 @@ dependencies = [ name = "openshell-driver-kubernetes" version = "0.0.0" dependencies = [ + "base64 0.22.1", "clap", "futures", "k8s-openapi", @@ -3943,10 +3944,12 @@ dependencies = [ "prost-types", "serde", "serde_json", + "sha2 0.10.9", "temp-env", "thiserror 2.0.18", "tokio", "tokio-stream", + "toml", "tonic", "tracing", "tracing-subscriber", diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index f211f906d6..91c20b0b22 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -64,7 +64,7 @@ delete, reconciliation removes the row; otherwise it can remain `Deleting`. |---|---|---|---| | Docker | Local development with Docker available. | Container plus nested sandbox namespace. | Uses host networking so loopback gateway endpoints work from the supervisor. | | Podman | Rootless or single-machine deployments. | Container plus nested sandbox namespace. | Uses the Podman REST API, OCI image volumes, and CDI GPU devices when available. | -| Kubernetes | Cluster deployment through Helm. | Pod plus nested sandbox namespace. | Uses Kubernetes API objects, service accounts, secrets, PVC-backed workspace storage, and GPU resources. | +| Kubernetes | Cluster deployment through Helm. | Pod plus nested sandbox namespace. | Uses Kubernetes API objects, service accounts, secrets, PVC-backed workspace storage, GPU resources, and optionally Agent Sandbox v1beta1 warm-pool claims. | | VM | Experimental microVM isolation. | Per-sandbox libkrun VM. | Managed endpoint-backed driver. The gateway spawns `openshell-driver-vm`, waits for its Unix socket, and then consumes it through the same remote `compute_driver.proto` path used by unmanaged endpoint drivers. The VM driver boots a cached bootstrap `rootfs.ext4`, prepares requested OCI images inside a bootstrap VM with `umoci`, attaches the prepared image disk read-only, and gives each sandbox a writable `overlay.ext4` for merged-root changes and runtime material. The driver persists each accepted launch request beside the overlay and restarts those VMs on driver startup without recreating the overlay. | | Extension | Out-of-tree drivers operated alongside the gateway. | Whatever boundary the driver implements. | Selected by a non-reserved custom `compute_drivers = [""]` entry with `[openshell.drivers.].socket_path`, or at launch time by pairing `--drivers ` with `--compute-driver-socket=`. Reserved built-in names such as `vm`, `docker`, `podman`, and `kubernetes` cannot be used as unmanaged socket endpoints. The gateway connects to a UDS the operator already provisioned, runs `GetCapabilities`, logs the advertised `driver_name`, and dispatches all sandbox lifecycle calls through `compute_driver.proto`. The driver process and socket lifecycle are operator-owned; the gateway does not spawn, supervise, or remove unmanaged extension drivers. The trust boundary is the socket's filesystem permissions: the operator must ensure only the gateway uid can read/write it. | @@ -125,6 +125,35 @@ Driver-controlled environment variables must override sandbox image or template values for sandbox ID, sandbox name, gateway endpoint, relay socket path, TLS paths, and command metadata. +Kubernetes supervisors authenticate to the gateway in two stages. They first +call `RegisterSupervisorPod` with the projected ServiceAccount token; the +gateway validates the pod-bound token and live Agent Sandbox owner state, then +activates already-bound cold pods by streaming back a gateway-minted sandbox +JWT. The supervisor installs that JWT in memory and starts the normal +`ConnectSupervisor` session as the activated sandbox. The gateway does not dial +pod IPs or require an inbound activation port; activation is supervisor +initiated over the existing outbound gRPC connection. + +When compatible OpenShell-enabled Agent Sandbox warm pools exist, the +Kubernetes driver can create a `SandboxClaim` instead of a direct `Sandbox`. +The driver returns an opaque activation token with the create response for +those delayed warm allocations. The gateway stores only a hash of that token +against the sandbox ID and later requires the claim activation controller to +present the same token before minting a sandbox-scoped supervisor JWT. The +token is derived by the driver from immutable live claim identity, including +the claim UID; Kubernetes labels remain useful for diagnostics and cross-checks +but are not sufficient activation proof. The driver watches warm pools and +templates into a local cache so create requests do not list or fetch extension +objects on the hot path. This remains behind the compute-driver boundary: +other drivers and the driver-agnostic gateway lifecycle continue to operate on +OpenShell sandbox identity and status. +The Kubernetes driver can optionally run a separate ConfigMap profile +reconciler that turns admin-authored TOML profiles into generated +`SandboxTemplate` and `SandboxWarmPool` resources. That reconciler is +independent of the create-path matcher: admins can use it, Helm, kubectl, or +another controller to create warm pools, and the matcher only considers the +resulting enabled warm-pool resources and their fingerprints. + Kubernetes can run the supervisor in the default combined topology or in a sidecar topology. Combined mode keeps network and process supervision in the agent container. Sidecar mode runs network enforcement, the proxy, and gateway diff --git a/architecture/gateway.md b/architecture/gateway.md index c8b323ea13..5b58798401 100644 --- a/architecture/gateway.md +++ b/architecture/gateway.md @@ -155,24 +155,36 @@ token is reported as connected but unauthenticated. Sandbox supervisor RPCs authenticate with explicit sandbox credentials; mTLS does not grant sandbox identity. Kubernetes deployments use the gateway-minted JWT bootstrap path: the supervisor starts with a projected -ServiceAccount token, exchanges it for a gateway-minted sandbox JWT, and uses -that JWT on subsequent gateway RPCs. +ServiceAccount token, registers the pod with the gateway, receives a +gateway-minted sandbox JWT after activation, and uses that JWT on subsequent +gateway RPCs. User-facing mutations are authorized by role policy when OIDC or edge identity is enabled. Sandbox secrets are gateway-signed JWTs bound to a single sandbox ID. Docker, Podman, and VM drivers deliver the initial token through supervisor-only runtime material; Kubernetes supervisors exchange a projected ServiceAccount -token through `IssueSandboxToken`. The gateway validates that projected token -with Kubernetes `TokenReview`, requires the configured sandbox service account, -checks the returned pod binding against the live pod UID, and verifies the pod's -controlling `Sandbox` ownerReference against the live Sandbox CR UID and -sandbox-id label before minting the gateway JWT. The bootstrap path accepts -both `agents.x-k8s.io/v1beta1` ownerReferences from newer Agent Sandbox -controllers and `agents.x-k8s.io/v1alpha1` ownerReferences from existing -deployments. Supervisors renew gateway JWTs in memory before expiry only while -the sandbox record still exists. Older tokens are not server-revoked; shared -deployments bound replay exposure with short `gateway_jwt.ttl_secs` lifetimes. +token through `RegisterSupervisorPod`. The gateway validates that projected +token with Kubernetes `TokenReview`, requires the configured sandbox service +account, and checks the returned pod binding against the live pod UID. Warm +pods that are not yet bound receive only a pending bootstrap identity and keep +the registration stream open. Later warm activation must match that pending +pod identity, the live Sandbox owner UID, and a gateway-stored opaque +activation token hash returned by the driver during sandbox creation. A missing +token mapping is retryable and leaves the pending stream intact; an empty or +mismatched token is fatal. Kubernetes sandbox-id labels are cross-checks and +diagnostics only, not authority to choose which sandbox receives the gateway +JWT. Already-bound pods still verify the pod's controlling `Sandbox` +ownerReference against the live Sandbox CR UID and sandbox-id label before +sending an activation message with the gateway JWT. +The bootstrap path accepts both `agents.x-k8s.io/v1beta1` ownerReferences from +newer Agent Sandbox controllers and `agents.x-k8s.io/v1alpha1` ownerReferences +from existing deployments. `IssueSandboxToken` remains as a compatibility shim +for older supervisor images and mints through the same already-bound pod +activation logic. Supervisors renew gateway JWTs in memory before expiry only +while the sandbox record still exists. Older tokens are not server-revoked; +shared deployments bound replay exposure with short `gateway_jwt.ttl_secs` +lifetimes. The config default is `gateway_jwt.ttl_secs = 0` for local single-player Docker, Podman, and VM gateways; those tokens carry `exp = 0` and do not expire. Kubernetes and other @@ -271,13 +283,16 @@ default WAL journal mode), which mirror the same sensitive contents. Persisted state includes sandboxes, providers, provider credential refresh state, SSH sessions, policy revisions, settings, inference configuration, and -deployment records. Provider refresh material is stored as a separate object -scoped to the provider instance through `objects.scope`; the provider record -keeps only the current injectable credential values and optional per-credential -expiry timestamps. A refresh normally mints one credential, but a strategy may -co-mint several (AWS STS mints the access key, secret key, and session token in -one call); the refresh state pins the resolved set of env keys it owns so -collision checks reserve all of them before the first mint. +deployment records. Warm-pool activation token mappings are stored as +generic objects keyed by sandbox ID with only the token hash in the payload; +sandbox deletion removes the mapping opportunistically after the sandbox row is +gone. Provider refresh material is stored as a separate object scoped to the +provider instance through `objects.scope`; the provider record keeps only the +current injectable credential values and optional per-credential expiry +timestamps. A refresh normally mints one credential, but a strategy may co-mint +several (AWS STS mints the access key, secret key, and session token in one +call); the refresh state pins the resolved set of env keys it owns so collision +checks reserve all of them before the first mint. ### Optimistic Concurrency (CAS) diff --git a/crates/openshell-cli/tests/ensure_providers_integration.rs b/crates/openshell-cli/tests/ensure_providers_integration.rs index 883c8c4446..3e56df00c2 100644 --- a/crates/openshell-cli/tests/ensure_providers_integration.rs +++ b/crates/openshell-cli/tests/ensure_providers_integration.rs @@ -564,6 +564,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-cli/tests/mtls_integration.rs b/crates/openshell-cli/tests/mtls_integration.rs index 622e3c1170..b5da73ecbb 100644 --- a/crates/openshell-cli/tests/mtls_integration.rs +++ b/crates/openshell-cli/tests/mtls_integration.rs @@ -449,6 +449,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-cli/tests/provider_commands_integration.rs b/crates/openshell-cli/tests/provider_commands_integration.rs index 53304b57c5..94e102c261 100644 --- a/crates/openshell-cli/tests/provider_commands_integration.rs +++ b/crates/openshell-cli/tests/provider_commands_integration.rs @@ -967,6 +967,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs b/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs index fda3d529f3..046f54c3d4 100644 --- a/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs +++ b/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs @@ -637,6 +637,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-cli/tests/sandbox_name_fallback_integration.rs b/crates/openshell-cli/tests/sandbox_name_fallback_integration.rs index 699749f0d1..8e8e65dbed 100644 --- a/crates/openshell-cli/tests/sandbox_name_fallback_integration.rs +++ b/crates/openshell-cli/tests/sandbox_name_fallback_integration.rs @@ -537,6 +537,17 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = tokio_stream::wrappers::ReceiverStream< + Result, + >; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-core/src/grpc_client.rs b/crates/openshell-core/src/grpc_client.rs index 070704fb0a..c261f3eee7 100644 --- a/crates/openshell-core/src/grpc_client.rs +++ b/crates/openshell-core/src/grpc_client.rs @@ -12,8 +12,8 @@ //! Podman / VM drivers write this to a bundle file at sandbox-create //! time). //! 3. `OPENSHELL_K8S_SA_TOKEN_FILE` — projected `ServiceAccount` JWT; the -//! supervisor exchanges it for a gateway JWT via `IssueSandboxToken` -//! once at startup. +//! supervisor registers its pod and receives a gateway JWT via +//! `RegisterSupervisorPod` once at startup. //! //! The resolved bearer credential is held in process memory thereafter and //! injected on every outbound call by [`AuthInterceptor`]. @@ -24,11 +24,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use crate::proto::{ DenialSummary, GetDraftPolicyRequest, GetInferenceBundleRequest, GetInferenceBundleResponse, - GetSandboxConfigRequest, GetSandboxProviderEnvironmentRequest, IssueSandboxTokenRequest, - NetworkActivitySummary, PolicyChunk, PolicySource, PolicyStatus, RefreshSandboxTokenRequest, - ReportPolicyStatusRequest, SandboxPolicy as ProtoSandboxPolicy, SubmitPolicyAnalysisRequest, - SubmitPolicyAnalysisResponse, UpdateConfigRequest, inference_client::InferenceClient, - open_shell_client::OpenShellClient, + GetSandboxConfigRequest, GetSandboxProviderEnvironmentRequest, NetworkActivitySummary, + PodActivationMessage, PolicyChunk, PolicySource, PolicyStatus, RefreshSandboxTokenRequest, + RegisterSupervisorPodRequest, ReportPolicyStatusRequest, SandboxPolicy as ProtoSandboxPolicy, + SubmitPolicyAnalysisRequest, SubmitPolicyAnalysisResponse, UpdateConfigRequest, + inference_client::InferenceClient, open_shell_client::OpenShellClient, }; use crate::sandbox_env; use miette::{IntoDiagnostic, Result, WrapErr}; @@ -242,7 +242,7 @@ async fn token_slot(endpoint: &str, plain_channel: &Channel) -> Result<(TokenSlo /// /// `endpoint` is logged on errors but never used for transport here; the /// actual network call lives inside this function only on the K8s -/// bootstrap path, which uses `plain_channel` to call `IssueSandboxToken` +/// bootstrap path, which uses `plain_channel` to call `RegisterSupervisorPod` /// once before the steady-state Bearer-authenticated channel is built. async fn acquire_sandbox_token(endpoint: &str, plain_channel: &Channel) -> Result { if let Ok(t) = std::env::var(sandbox_env::SANDBOX_TOKEN) @@ -272,7 +272,9 @@ async fn acquire_sandbox_token(endpoint: &str, plain_channel: &Channel) -> Resul && !sa_path.is_empty() { return Ok(AcquiredToken { - token: acquire_k8s_sandbox_token(endpoint, plain_channel, &sa_path).await?, + token: acquire_k8s_supervisor_activation(endpoint, plain_channel, &sa_path) + .await? + .token, refresh_mode: RefreshMode::GatewayJwt(TokenSource::K8sServiceAccount), }); } @@ -290,12 +292,24 @@ async fn acquire_k8s_sandbox_token( plain_channel: &Channel, sa_path: &str, ) -> Result { + Ok( + acquire_k8s_supervisor_activation(endpoint, plain_channel, sa_path) + .await? + .token, + ) +} + +async fn acquire_k8s_supervisor_activation( + endpoint: &str, + plain_channel: &Channel, + sa_path: &str, +) -> Result { let sa_token = std::fs::read_to_string(sa_path) .into_diagnostic() .wrap_err_with(|| format!("failed to read K8s SA token from {sa_path}"))? .trim() .to_string(); - info!(endpoint = %endpoint, "exchanging K8s ServiceAccount token for sandbox JWT"); + info!(endpoint = %endpoint, "registering K8s supervisor pod for sandbox activation"); // The bootstrap exchange uses a one-off interceptor pinned to the // SA token; the resulting gateway JWT becomes the value in the // shared `TOKEN_SLOT` once `connect_channel` returns. @@ -308,11 +322,77 @@ async fn acquire_k8s_sandbox_token( let bootstrap = InterceptedService::new(plain_channel.clone(), interceptor); let mut client = OpenShellClient::new(bootstrap); let resp = client - .issue_sandbox_token(IssueSandboxTokenRequest {}) + .register_supervisor_pod(RegisterSupervisorPodRequest {}) + .await + .into_diagnostic() + .wrap_err("RegisterSupervisorPod bootstrap stream failed")?; + let mut stream = resp.into_inner(); + let activation = stream + .message() .await .into_diagnostic() - .wrap_err("IssueSandboxToken bootstrap exchange failed")?; - Ok(resp.into_inner().token) + .wrap_err("RegisterSupervisorPod activation stream failed")? + .ok_or_else(|| miette::miette!("RegisterSupervisorPod stream closed before activation"))?; + info!( + sandbox_id = %activation.sandbox_id, + sandbox_name = %activation.sandbox_name, + "received supervisor pod activation" + ); + Ok(activation) +} + +/// Register a Kubernetes supervisor pod and install the activated gateway JWT. +/// +/// Warm pods call this before policy loading. The registration stream remains +/// pending until the Kubernetes driver observes a claim binding and the gateway +/// sends activation. +pub async fn register_k8s_supervisor_pod(endpoint: &str) -> Result { + let sa_path = std::env::var(sandbox_env::K8S_SA_TOKEN_FILE) + .ok() + .filter(|path| !path.is_empty()) + .ok_or_else(|| { + miette::miette!( + "{} must be set to register Kubernetes supervisor pod", + sandbox_env::K8S_SA_TOKEN_FILE + ) + })?; + + let mut backoff = Duration::from_secs(1); + loop { + let activation = { + let guard = TOKEN_INIT_LOCK.lock().await; + if TOKEN_SLOT.get().is_some() { + return Err(miette::miette!( + "Kubernetes supervisor pod registration was requested after sandbox token initialization" + )); + } + + match async { + let plain_channel = build_plain_channel(endpoint).await?; + acquire_k8s_supervisor_activation(endpoint, &plain_channel, &sa_path).await + } + .await + { + Ok(activation) => activation, + Err(err) => { + warn!( + endpoint = %endpoint, + error = %err, + retry_after_secs = backoff.as_secs(), + "Kubernetes supervisor pod registration failed; retrying: {err:#}" + ); + drop(guard); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(Duration::from_secs(30)); + continue; + } + } + }; + + let _slot = install_token_slot(&activation.token)?; + let _ = TOKEN_REFRESH_MODE.set(RefreshMode::GatewayJwt(TokenSource::K8sServiceAccount)); + return Ok(activation); + } } /// Build an authenticated channel for direct external use (e.g. the diff --git a/crates/openshell-core/src/lib.rs b/crates/openshell-core/src/lib.rs index 20c4dccb7d..f007fc178c 100644 --- a/crates/openshell-core/src/lib.rs +++ b/crates/openshell-core/src/lib.rs @@ -37,6 +37,7 @@ pub mod provider_credentials; pub mod sandbox_env; pub mod secrets; pub mod settings; +pub mod supervisor_bootstrap; pub mod telemetry; pub mod time; diff --git a/crates/openshell-core/src/sandbox_env.rs b/crates/openshell-core/src/sandbox_env.rs index f066580636..08da8b946c 100644 --- a/crates/openshell-core/src/sandbox_env.rs +++ b/crates/openshell-core/src/sandbox_env.rs @@ -87,9 +87,10 @@ pub const USER_ENVIRONMENT: &str = "OPENSHELL_USER_ENVIRONMENT"; /// Path to the projected `ServiceAccount` JWT (Kubernetes driver). /// -/// Used to bootstrap a gateway-minted JWT via `IssueSandboxToken`. Kubelet -/// writes and rotates this file; the supervisor exchanges its contents -/// for a gateway JWT at startup and on refresh. +/// Used to register the supervisor pod and receive a gateway-minted JWT via +/// `RegisterSupervisorPod`. Kubelet writes and rotates this file; the +/// supervisor presents its contents at startup and when rebootstrap is needed +/// after refresh authentication failure. pub const K8S_SA_TOKEN_FILE: &str = "OPENSHELL_K8S_SA_TOKEN_FILE"; /// Filesystem path to the SPIFFE Workload API UNIX socket used for provider diff --git a/crates/openshell-core/src/supervisor_bootstrap.rs b/crates/openshell-core/src/supervisor_bootstrap.rs new file mode 100644 index 0000000000..bda04d12e9 --- /dev/null +++ b/crates/openshell-core/src/supervisor_bootstrap.rs @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Driver-provided supervisor bootstrap capabilities. +//! +//! These types describe the narrow interface between the gateway and compute +//! drivers for supervisors that cannot start with a gateway-minted sandbox JWT. +//! Kubernetes is the initial implementation: the driver validates projected +//! `ServiceAccount` tokens and classifies pods as already bound or warm-pending. + +use tonic::Status; +use tonic::async_trait; + +/// Registration-only identity for a supervisor runtime instance. +/// +/// This is not a sandbox principal. The gateway may use it only on the +/// bootstrap registration path until a concrete sandbox token is minted. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SupervisorBootstrapIdentity { + /// Driver that produced the identity, for example `kubernetes`. + pub driver: String, + /// Driver-native stable instance ID, for example Kubernetes pod UID. + pub instance_id: String, + /// Driver-native human-readable instance name, for example Kubernetes pod + /// name. + pub instance_name: String, + /// Driver-native owner object name. + pub owner_name: String, + /// Driver-native owner object UID. + pub owner_uid: String, + /// Whether the instance is already bound to an `OpenShell` sandbox or is + /// waiting for a warm-pool claim. + pub binding: SupervisorBootstrapBinding, +} + +/// Binding state returned by a driver bootstrap identity provider. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum SupervisorBootstrapBinding { + /// The instance is already bound to a concrete `OpenShell` sandbox. + BoundSandbox { sandbox_id: String }, + /// The instance is valid but not yet claim-bound. + WarmPending, +} + +impl SupervisorBootstrapIdentity { + /// Return the bound sandbox ID when the identity is already activated by + /// driver state. + #[must_use] + pub fn bound_sandbox_id(&self) -> Option<&str> { + match &self.binding { + SupervisorBootstrapBinding::BoundSandbox { sandbox_id } => Some(sandbox_id.as_str()), + SupervisorBootstrapBinding::WarmPending => None, + } + } +} + +/// Driver-provided authentication for supervisor bootstrap registration. +#[async_trait] +pub trait SupervisorBootstrapIdentityProvider: Send + Sync { + /// Authenticate a driver-native bootstrap token. + /// + /// `Ok(None)` means the token did not authenticate and another + /// authenticator may try it. `Err` means authentication could not safely + /// complete or the token was authenticated but invalid for bootstrap. + async fn authenticate_registration( + &self, + token: &str, + ) -> Result, Status>; +} + +/// Request sent by a driver-side warm-pool controller to activate a pending +/// bootstrap stream. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SupervisorBootstrapActivationRequest { + /// Driver that produced the pending registration. + pub driver: String, + /// Driver-native stable instance ID to activate. + pub instance_id: String, + /// `OpenShell` sandbox ID the instance should become. + pub sandbox_id: String, + /// Driver-native owner UID observed during claim validation. + pub owner_uid: String, + /// Opaque driver-issued activation token returned during sandbox create. + pub activation_token: String, + /// Short reason for logs and audit messages. + pub reason: String, +} + +/// Gateway-owned activation callback passed to driver-side warm-pool logic. +#[async_trait] +pub trait SupervisorBootstrapActivator: Send + Sync { + async fn activate_registered_supervisor( + &self, + request: SupervisorBootstrapActivationRequest, + ) -> Result<(), Status>; +} diff --git a/crates/openshell-driver-docker/src/lib.rs b/crates/openshell-driver-docker/src/lib.rs index 2f89c2229e..f71957e703 100644 --- a/crates/openshell-driver-docker/src/lib.rs +++ b/crates/openshell-driver-docker/src/lib.rs @@ -1430,7 +1430,9 @@ impl ComputeDriver for DockerComputeDriver { .sandbox .ok_or_else(|| Status::invalid_argument("sandbox is required"))?; self.create_sandbox_inner(&sandbox).await?; - Ok(Response::new(CreateSandboxResponse {})) + Ok(Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( diff --git a/crates/openshell-driver-kubernetes/Cargo.toml b/crates/openshell-driver-kubernetes/Cargo.toml index 2c02f864ab..983e366923 100644 --- a/crates/openshell-driver-kubernetes/Cargo.toml +++ b/crates/openshell-driver-kubernetes/Cargo.toml @@ -29,6 +29,9 @@ kube-runtime = { workspace = true } k8s-openapi = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +toml = { workspace = true } +sha2 = { workspace = true } +base64 = { workspace = true } clap = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/crates/openshell-driver-kubernetes/README.md b/crates/openshell-driver-kubernetes/README.md index 96e54ad448..9af04155d0 100644 --- a/crates/openshell-driver-kubernetes/README.md +++ b/crates/openshell-driver-kubernetes/README.md @@ -26,6 +26,35 @@ by the gateway. Kubernetes API calls use explicit timeouts so gRPC handlers do not block indefinitely when the API server is slow or unavailable. +## Warm-Pool Allocation + +The driver can use Agent Sandbox extension CRDs for transparent warm-pool +allocation. When `warm_pooling.enabled` is true, the driver watches and caches +`extensions.agents.x-k8s.io/v1beta1` `SandboxWarmPool` resources labelled +`openshell.ai/enabled=true`, resolves each pool's referenced +`SandboxTemplate`, and computes an in-memory fingerprint of the template spec. + +On create, the driver renders the Kubernetes spec that direct `Sandbox` +creation would use, strips per-sandbox identity values from the fingerprint, +and matches it against the warm-pool cache for the target Kubernetes namespace. +If exactly one pool matches, the driver creates a v1beta1 `SandboxClaim` with +`spec.warmPoolRef.name` set to that pool. If no pool matches, multiple pools +match, RBAC prevents cache maintenance, or the v1beta1 extension APIs are +unavailable, the driver falls back to direct `Sandbox` creation. + +The extension CRDs are intentionally v1beta1-only. The core +`agents.x-k8s.io/Sandbox` CRD still supports the existing v1beta1-to-v1alpha1 +fallback. + +The driver can also reconcile admin-authored warm-pool profiles from labelled +ConfigMaps when `warm_pooling.profiles.enabled` is true. Profiles live in the +configured profile namespace, use the `warm-pool.toml` data key, and expose a +narrow CLI-shaped schema: `workspace`, `replicas`, `image`, +`runtime_class_name`, `[environment]`, and `[resources]` with `cpu`, `memory`, +and `gpu_count`. The reconciler generates ordinary `SandboxTemplate` and +`SandboxWarmPool` resources labelled `openshell.ai/enabled=true`; the existing +matching cache then handles allocation unchanged. + ## Workspace Persistence Sandbox pods use a PVC-backed `/sandbox` workspace. An init container seeds the @@ -45,8 +74,9 @@ values must override image-provided environment variables. Sandbox pods run as `service_account_name` and keep `automountServiceAccountToken: false`. The only Kubernetes token exposed to the supervisor is an explicit, audience-bound projected token mounted at -`/var/run/secrets/openshell/token` for the one-shot `IssueSandboxToken` -bootstrap exchange. +`/var/run/secrets/openshell/token` for the `RegisterSupervisorPod` bootstrap +stream. The gateway validates the pod-bound token and activates already-bound +cold pods by returning a gateway-minted sandbox JWT on that stream. The gateway uses the supervisor relay for connect, exec, and file sync. Sandbox pods do not need direct external ingress for SSH. diff --git a/crates/openshell-driver-kubernetes/src/bootstrap.rs b/crates/openshell-driver-kubernetes/src/bootstrap.rs new file mode 100644 index 0000000000..463339adc7 --- /dev/null +++ b/crates/openshell-driver-kubernetes/src/bootstrap.rs @@ -0,0 +1,907 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Kubernetes `ServiceAccount` supervisor bootstrap identity provider. +//! +//! Validates a projected SA token presented by a sandbox pod, reads the pod's +//! `openshell.io/sandbox-id` annotation, verifies the pod is controlled by the +//! corresponding Sandbox CR, and returns a registration-only driver bootstrap +//! identity. Warm pods may register before they are bound to a sandbox, so this +//! identity must not grant sandbox-scoped RPC access. +//! +//! This is the Kubernetes driver's apiserver-facing side of the supervisor +//! bootstrap boundary. The gateway owns the public registration stream and +//! token minting. + +use k8s_openapi::api::{ + authentication::v1::{TokenReview, TokenReviewSpec, TokenReviewStatus, UserInfo}, + core::v1::Pod, +}; +use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta; +use kube::Error as KubeError; +use kube::api::{Api, ApiResource, PostParams}; +use kube::core::{DynamicObject, gvk::GroupVersionKind}; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapBinding, SupervisorBootstrapIdentity, SupervisorBootstrapIdentityProvider, +}; +use std::sync::Arc; +use tonic::Status; +use tonic::async_trait; +use tracing::{debug, info, warn}; + +/// Pod annotation that binds a sandbox pod to its UUID. Set by the +/// Kubernetes compute driver at pod-create time. +pub const SANDBOX_ID_ANNOTATION: &str = "openshell.io/sandbox-id"; +const SANDBOX_API_GROUP: &str = "agents.x-k8s.io"; +const SANDBOX_API_VERSION_V1BETA1: &str = "v1beta1"; +const SANDBOX_API_VERSION_V1ALPHA1: &str = "v1alpha1"; +const SANDBOX_API_VERSION_FULL_V1BETA1: &str = "agents.x-k8s.io/v1beta1"; +const SANDBOX_API_VERSION_FULL_V1ALPHA1: &str = "agents.x-k8s.io/v1alpha1"; +const SANDBOX_KIND: &str = "Sandbox"; +const SANDBOX_ID_LABEL: &str = "openshell.ai/sandbox-id"; +const POD_NAME_EXTRA: &str = "authentication.kubernetes.io/pod-name"; +const POD_UID_EXTRA: &str = "authentication.kubernetes.io/pod-uid"; + +/// Apiserver-facing operations the authenticator depends on. Split out so +/// tests can fake the apiserver without standing up a kube cluster. +#[async_trait] +pub trait K8sIdentityResolver: Send + Sync + 'static { + /// Validate `token` via `TokenReview` (`aud == openshell-gateway`), + /// extract the pod name/uid, then `GET` the pod and owning Sandbox CR. + /// Returns `Ok(None)` when the token is well-formed but does not + /// authenticate (e.g. wrong audience); returns `Err` for + /// transport/server errors. + async fn resolve(&self, token: &str) -> Result, Status>; +} + +/// Kubernetes implementation of the driver bootstrap identity provider. +pub struct KubernetesSupervisorBootstrapIdentityProvider { + resolver: Arc, +} + +impl std::fmt::Debug for KubernetesSupervisorBootstrapIdentityProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("KubernetesSupervisorBootstrapIdentityProvider") + .finish_non_exhaustive() + } +} + +impl KubernetesSupervisorBootstrapIdentityProvider { + pub fn new(resolver: Arc) -> Self { + Self { resolver } + } +} + +#[async_trait] +impl SupervisorBootstrapIdentityProvider for KubernetesSupervisorBootstrapIdentityProvider { + async fn authenticate_registration( + &self, + token: &str, + ) -> Result, Status> { + self.resolver.resolve(token).await + } +} + +#[derive(Debug)] +struct TokenReviewIdentity { + pod_name: String, + pod_uid: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct SandboxOwnerReference { + api_version: String, + name: String, + uid: String, +} + +/// Resolver backed by the apiserver's `TokenReview` API and `kube::Client` +/// for the per-pod annotation lookup. +pub struct LiveK8sResolver { + token_reviews_api: Api, + pods_api: Api, + sandboxes_api_v1beta1: Api, + sandboxes_api_v1alpha1: Api, + expected_audience: String, + sandbox_namespace: String, + expected_service_account: String, +} + +impl LiveK8sResolver { + pub fn new( + client: kube::Client, + namespace: &str, + expected_audience: String, + expected_service_account: String, + ) -> Self { + let token_reviews_api: Api = Api::all(client.clone()); + let pods_api: Api = Api::namespaced(client.clone(), namespace); + let sandbox_gvk_v1beta1 = + GroupVersionKind::gvk(SANDBOX_API_GROUP, SANDBOX_API_VERSION_V1BETA1, SANDBOX_KIND); + let sandbox_resource_v1beta1 = ApiResource::from_gvk(&sandbox_gvk_v1beta1); + let sandbox_gvk_v1alpha1 = GroupVersionKind::gvk( + SANDBOX_API_GROUP, + SANDBOX_API_VERSION_V1ALPHA1, + SANDBOX_KIND, + ); + let sandbox_resource_v1alpha1 = ApiResource::from_gvk(&sandbox_gvk_v1alpha1); + let sandboxes_api_v1beta1: Api = + Api::namespaced_with(client.clone(), namespace, &sandbox_resource_v1beta1); + let sandboxes_api_v1alpha1: Api = + Api::namespaced_with(client, namespace, &sandbox_resource_v1alpha1); + Self { + token_reviews_api, + pods_api, + sandboxes_api_v1beta1, + sandboxes_api_v1alpha1, + expected_audience, + sandbox_namespace: namespace.to_string(), + expected_service_account, + } + } + + async fn get_sandbox_cr_for_owner( + &self, + owner: &SandboxOwnerReference, + ) -> Result, KubeError> { + let apis = if owner.api_version == SANDBOX_API_VERSION_FULL_V1ALPHA1 { + [&self.sandboxes_api_v1alpha1, &self.sandboxes_api_v1beta1] + } else { + [&self.sandboxes_api_v1beta1, &self.sandboxes_api_v1alpha1] + }; + + for api in apis { + match api.get_opt(&owner.name).await { + Ok(Some(sandbox_cr)) => return Ok(Some(sandbox_cr)), + Ok(None) => {} + Err(err) if should_try_next_sandbox_api_version(&err) => {} + Err(err) => return Err(err), + } + } + + Ok(None) + } +} + +#[async_trait] +impl K8sIdentityResolver for LiveK8sResolver { + async fn resolve(&self, token: &str) -> Result, Status> { + let review = TokenReview { + metadata: ObjectMeta::default(), + spec: TokenReviewSpec { + audiences: Some(vec![self.expected_audience.clone()]), + token: Some(token.to_string()), + }, + status: None, + }; + + let review = self + .token_reviews_api + .create(&PostParams::default(), &review) + .await + .map_err(|e| { + warn!(error = %e, "K8s TokenReview failed"); + Status::internal(format!("tokenreview failed: {e}")) + })?; + let status = review + .status + .ok_or_else(|| Status::internal("TokenReview response missing status"))?; + let Some(identity) = token_review_identity( + &status, + &self.expected_audience, + &self.sandbox_namespace, + &self.expected_service_account, + )? + else { + return Ok(None); + }; + + info!( + pod_name = %identity.pod_name, + pod_uid = %identity.pod_uid, + service_account = %self.expected_service_account, + "validated K8s SA token via TokenReview" + ); + + // Look up the pod and read its sandbox-id annotation. + let pod = self + .pods_api + .get_opt(&identity.pod_name) + .await + .map_err(|e| { + warn!( + pod = %identity.pod_name, + error = %e, + "failed to fetch sandbox pod for annotation lookup" + ); + Status::internal(format!("pod GET failed: {e}")) + })?; + let Some(pod) = pod else { + warn!( + pod = %identity.pod_name, + "sandbox pod referenced by SA token not found in this namespace" + ); + return Err(Status::not_found("sandbox pod not found")); + }; + + // Defense-in-depth: confirm the pod UID matches the SA token's + // `kubernetes.io.pod.uid`. Prevents a replayed token from a + // recreated pod with the same name. + let actual_uid = pod.metadata.uid.as_deref().unwrap_or_default(); + if actual_uid != identity.pod_uid { + warn!( + pod = %identity.pod_name, + claimed_uid = %identity.pod_uid, + actual_uid = %actual_uid, + "SA token pod UID does not match live pod; rejecting" + ); + return Err(Status::permission_denied("SA token pod UID mismatch")); + } + + let sandbox_id = pod_sandbox_id(&pod); + + let owner = sandbox_owner_reference(&pod)?; + let sandbox_cr = self.get_sandbox_cr_for_owner(&owner).await.map_err(|e| { + warn!( + pod = %identity.pod_name, + sandbox_owner = %owner.name, + sandbox_owner_api_version = %owner.api_version, + error = %e, + "failed to fetch owning Sandbox CR for pod identity validation" + ); + Status::internal(format!("sandbox GET failed: {e}")) + })?; + let Some(sandbox_cr) = sandbox_cr else { + warn!( + pod = %identity.pod_name, + sandbox_owner = %owner.name, + sandbox_owner_api_version = %owner.api_version, + "pod ownerReference points to a Sandbox CR that does not exist" + ); + return Err(Status::permission_denied("sandbox owner not found")); + }; + validate_sandbox_owner_reference(&owner, &sandbox_cr)?; + if let Some(ref sandbox_id) = sandbox_id { + validate_sandbox_owner_binding(&owner, sandbox_id, &sandbox_cr)?; + } + + let binding = sandbox_id.map_or(SupervisorBootstrapBinding::WarmPending, |sandbox_id| { + SupervisorBootstrapBinding::BoundSandbox { sandbox_id } + }); + + Ok(Some(SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: identity.pod_name, + instance_id: identity.pod_uid, + owner_name: owner.name, + owner_uid: owner.uid, + binding, + })) + } +} + +#[allow(clippy::result_large_err)] +fn token_review_identity( + status: &TokenReviewStatus, + expected_audience: &str, + sandbox_namespace: &str, + expected_service_account: &str, +) -> Result, Status> { + if status.authenticated != Some(true) { + debug!( + error = status.error.as_deref().unwrap_or_default(), + "K8s TokenReview did not authenticate token" + ); + return Ok(None); + } + + let audiences = status.audiences.as_deref().unwrap_or_default(); + if !audiences.iter().any(|aud| aud == expected_audience) { + warn!( + expected_audience = %expected_audience, + audiences = ?audiences, + "K8s TokenReview authenticated token without expected audience" + ); + return Err(Status::unauthenticated("SA token audience not accepted")); + } + + let user = status + .user + .as_ref() + .ok_or_else(|| Status::permission_denied("TokenReview response missing user info"))?; + let username = user + .username + .as_deref() + .ok_or_else(|| Status::permission_denied("TokenReview response missing username"))?; + let expected_username = + format!("system:serviceaccount:{sandbox_namespace}:{expected_service_account}"); + if username != expected_username { + warn!( + username = %username, + sandbox_namespace = %sandbox_namespace, + service_account = %expected_service_account, + "K8s TokenReview principal is not the configured sandbox service account" + ); + return Err(Status::permission_denied( + "SA token is not from the configured sandbox service account", + )); + } + + let pod_name = user_extra_one(user, POD_NAME_EXTRA)?; + let pod_uid = user_extra_one(user, POD_UID_EXTRA)?; + Ok(Some(TokenReviewIdentity { pod_name, pod_uid })) +} + +#[allow(clippy::result_large_err)] +fn user_extra_one(user: &UserInfo, key: &str) -> Result { + let Some(values) = user.extra.as_ref().and_then(|extra| extra.get(key)) else { + return Err(Status::permission_denied("SA token is not pod-bound")); + }; + if values.len() != 1 || values[0].is_empty() { + return Err(Status::permission_denied( + "SA token has invalid pod binding", + )); + } + Ok(values[0].clone()) +} + +#[allow(clippy::result_large_err)] +fn pod_sandbox_id(pod: &Pod) -> Option { + pod.metadata + .annotations + .as_ref() + .and_then(|a| a.get(SANDBOX_ID_ANNOTATION)) + .cloned() + .filter(|sandbox_id| !sandbox_id.is_empty()) +} + +#[allow(clippy::result_large_err)] +fn sandbox_owner_reference(pod: &Pod) -> Result { + let owner_refs = pod.metadata.owner_references.as_deref().unwrap_or_default(); + let mut sandbox_refs = owner_refs + .iter() + .filter(|owner| is_supported_sandbox_owner_reference(owner)); + let Some(owner) = sandbox_refs.next() else { + let unsupported_sandbox_api_versions = owner_refs + .iter() + .filter(|owner| owner.kind == SANDBOX_KIND) + .map(|owner| owner.api_version.as_str()) + .collect::>(); + if !unsupported_sandbox_api_versions.is_empty() { + warn!( + api_versions = ?unsupported_sandbox_api_versions, + supported_api_versions = ?[ + SANDBOX_API_VERSION_FULL_V1BETA1, + SANDBOX_API_VERSION_FULL_V1ALPHA1, + ], + "pod Sandbox ownerReference uses unsupported apiVersion" + ); + } + return Err(Status::permission_denied( + "pod is not controlled by an OpenShell Sandbox", + )); + }; + if sandbox_refs.next().is_some() { + return Err(Status::permission_denied( + "pod has multiple OpenShell Sandbox owners", + )); + } + if owner.controller != Some(true) { + return Err(Status::permission_denied( + "pod Sandbox ownerReference is not controlling", + )); + } + if owner.name.is_empty() || owner.uid.is_empty() { + return Err(Status::permission_denied( + "pod Sandbox ownerReference is incomplete", + )); + } + Ok(SandboxOwnerReference { + api_version: owner.api_version.clone(), + name: owner.name.clone(), + uid: owner.uid.clone(), + }) +} + +fn is_supported_sandbox_owner_reference( + owner: &k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference, +) -> bool { + owner.kind == SANDBOX_KIND + && matches!( + owner.api_version.as_str(), + SANDBOX_API_VERSION_FULL_V1BETA1 | SANDBOX_API_VERSION_FULL_V1ALPHA1 + ) +} + +fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { + // Kubernetes returns a structured 404 for some missing API resources and a + // raw "404 page not found" body for others. Both mean the probed + // group/version is unavailable and the next supported Sandbox API version + // should be tried. + matches!(err, KubeError::Api(api) if api.code == 404) +} + +#[allow(clippy::result_large_err)] +fn validate_sandbox_owner_reference( + owner: &SandboxOwnerReference, + sandbox_cr: &DynamicObject, +) -> Result<(), Status> { + let actual_uid = sandbox_cr.metadata.uid.as_deref().unwrap_or_default(); + if actual_uid != owner.uid { + warn!( + sandbox_owner = %owner.name, + owner_uid = %owner.uid, + actual_uid = %actual_uid, + "pod Sandbox ownerReference UID does not match live Sandbox CR" + ); + return Err(Status::permission_denied("sandbox owner UID mismatch")); + } + + Ok(()) +} + +#[allow(clippy::result_large_err)] +fn validate_sandbox_owner_binding( + owner: &SandboxOwnerReference, + sandbox_id: &str, + sandbox_cr: &DynamicObject, +) -> Result<(), Status> { + let actual_sandbox_id = sandbox_cr + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(SANDBOX_ID_LABEL)) + .map(String::as_str) + .unwrap_or_default(); + if actual_sandbox_id != sandbox_id { + warn!( + sandbox_owner = %owner.name, + owner_uid = %owner.uid, + pod_sandbox_id = %sandbox_id, + cr_sandbox_id = %actual_sandbox_id, + "pod sandbox annotation does not match owning Sandbox CR label" + ); + return Err(Status::permission_denied("sandbox owner ID mismatch")); + } + + Ok(()) +} + +#[cfg(test)] +pub mod test_support { + use super::*; + use std::sync::Mutex; + + /// Fake resolver for unit tests. Returns the configured outcome on + /// every call and records the tokens it observed. + pub struct FakeResolver { + pub outcome: Result, Status>, + pub seen_tokens: Mutex>, + } + + impl FakeResolver { + pub fn returning(outcome: Result, Status>) -> Self { + Self { + outcome, + seen_tokens: Mutex::new(Vec::new()), + } + } + } + + #[async_trait] + impl K8sIdentityResolver for FakeResolver { + async fn resolve( + &self, + token: &str, + ) -> Result, Status> { + self.seen_tokens.lock().unwrap().push(token.to_string()); + match &self.outcome { + Ok(opt) => Ok(opt.clone()), + Err(s) => Err(Status::new(s.code(), s.message())), + } + } + } +} + +#[cfg(test)] +mod tests { + use super::test_support::FakeResolver; + use super::*; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference; + use std::collections::BTreeMap; + use std::sync::Arc; + + fn kube_api_error(code: u16, message: &str) -> KubeError { + KubeError::Api(kube::core::ErrorResponse { + status: if code == 404 { + "404 Not Found".to_string() + } else { + "Failure".to_string() + }, + message: message.to_string(), + reason: "Failed to parse error data".to_string(), + code, + }) + } + + #[test] + fn sandbox_api_version_probe_retries_on_structured_and_raw_404() { + let structured = kube_api_error(404, "could not find the requested resource"); + assert!(should_try_next_sandbox_api_version(&structured)); + + let raw = kube_api_error(404, "404 page not found\n"); + assert!(should_try_next_sandbox_api_version(&raw)); + } + + #[test] + fn sandbox_api_version_probe_keeps_non_404_errors() { + let err = kube_api_error(403, "sandboxes.agents.x-k8s.io is forbidden"); + assert!(!should_try_next_sandbox_api_version(&err)); + } + + fn token_review_status( + authenticated: bool, + audiences: Vec<&str>, + username: &str, + extra: Vec<(&str, &str)>, + ) -> TokenReviewStatus { + TokenReviewStatus { + authenticated: Some(authenticated), + audiences: Some(audiences.into_iter().map(str::to_string).collect()), + error: None, + user: Some(UserInfo { + username: Some(username.to_string()), + uid: Some("sa-uid".to_string()), + groups: Some(vec![ + "system:serviceaccounts".to_string(), + "system:serviceaccounts:openshell".to_string(), + "system:authenticated".to_string(), + ]), + extra: Some( + extra + .into_iter() + .map(|(k, v)| (k.to_string(), vec![v.to_string()])) + .collect::>(), + ), + }), + } + } + + fn sandbox_owner(name: &str, uid: &str) -> OwnerReference { + sandbox_owner_with_api_version(SANDBOX_API_VERSION_FULL_V1BETA1, name, uid) + } + + fn sandbox_owner_with_api_version(api_version: &str, name: &str, uid: &str) -> OwnerReference { + OwnerReference { + api_version: api_version.to_string(), + block_owner_deletion: None, + controller: Some(true), + kind: SANDBOX_KIND.to_string(), + name: name.to_string(), + uid: uid.to_string(), + } + } + + fn pod_with_owner_refs(owner_references: Vec) -> Pod { + Pod { + metadata: ObjectMeta { + owner_references: Some(owner_references), + ..Default::default() + }, + ..Default::default() + } + } + + fn pod_with_sandbox_id(sandbox_id: Option<&str>) -> Pod { + Pod { + metadata: ObjectMeta { + annotations: sandbox_id.map(|id| { + BTreeMap::from([(SANDBOX_ID_ANNOTATION.to_string(), id.to_string())]) + }), + ..Default::default() + }, + ..Default::default() + } + } + + fn bootstrap_identity(binding: SupervisorBootstrapBinding) -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "openshell-sandbox-a".to_string(), + instance_id: "uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "cr-uid-a".to_string(), + binding, + } + } + + fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str) -> DynamicObject { + let sandbox_gvk = + GroupVersionKind::gvk(SANDBOX_API_GROUP, SANDBOX_API_VERSION_V1BETA1, SANDBOX_KIND); + let sandbox_resource = ApiResource::from_gvk(&sandbox_gvk); + let mut cr = DynamicObject::new(name, &sandbox_resource); + cr.metadata.uid = Some(uid.to_string()); + cr.metadata.labels = Some(BTreeMap::from([( + SANDBOX_ID_LABEL.to_string(), + sandbox_id.to_string(), + )])); + cr + } + + #[test] + fn token_review_identity_extracts_pod_binding() { + let status = token_review_status( + true, + vec!["openshell-gateway"], + "system:serviceaccount:openshell:default", + vec![ + (POD_NAME_EXTRA, "openshell-sandbox-a"), + (POD_UID_EXTRA, "uid-a"), + ], + ); + + let identity = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .unwrap() + .expect("authenticated token should resolve"); + + assert_eq!(identity.pod_name, "openshell-sandbox-a"); + assert_eq!(identity.pod_uid, "uid-a"); + } + + #[test] + fn token_review_identity_returns_none_when_not_authenticated() { + let status = TokenReviewStatus { + authenticated: Some(false), + error: Some("invalid audience".to_string()), + ..Default::default() + }; + + assert!( + token_review_identity(&status, "openshell-gateway", "openshell", "default") + .unwrap() + .is_none() + ); + } + + #[test] + fn token_review_identity_requires_expected_audience() { + let status = token_review_status( + true, + vec!["kubernetes.default.svc"], + "system:serviceaccount:openshell:default", + vec![ + (POD_NAME_EXTRA, "openshell-sandbox-a"), + (POD_UID_EXTRA, "uid-a"), + ], + ); + + let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .expect_err("wrong audience must fail closed"); + assert_eq!(err.code(), tonic::Code::Unauthenticated); + } + + #[test] + fn token_review_identity_requires_sandbox_namespace() { + let status = token_review_status( + true, + vec!["openshell-gateway"], + "system:serviceaccount:other:default", + vec![ + (POD_NAME_EXTRA, "openshell-sandbox-a"), + (POD_UID_EXTRA, "uid-a"), + ], + ); + + let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .expect_err("other namespace must be rejected"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn token_review_identity_requires_configured_service_account() { + let status = token_review_status( + true, + vec!["openshell-gateway"], + "system:serviceaccount:openshell:other", + vec![ + (POD_NAME_EXTRA, "openshell-sandbox-a"), + (POD_UID_EXTRA, "uid-a"), + ], + ); + + let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .expect_err("other service account must be rejected"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn token_review_identity_requires_pod_bound_extras() { + let status = token_review_status( + true, + vec!["openshell-gateway"], + "system:serviceaccount:openshell:default", + vec![], + ); + + let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") + .expect_err("non pod-bound tokens must be rejected"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn pod_sandbox_id_is_optional_for_warm_registration() { + assert_eq!( + pod_sandbox_id(&pod_with_sandbox_id(Some("sandbox-id-a"))).as_deref(), + Some("sandbox-id-a") + ); + assert!(pod_sandbox_id(&pod_with_sandbox_id(None)).is_none()); + } + + #[test] + fn sandbox_owner_reference_extracts_controlling_sandbox_owner() { + let pod = pod_with_owner_refs(vec![sandbox_owner("sandbox-a", "cr-uid-a")]); + + let owner = sandbox_owner_reference(&pod).expect("expected Sandbox owner"); + + assert_eq!( + owner, + SandboxOwnerReference { + api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), + name: "sandbox-a".to_string(), + uid: "cr-uid-a".to_string(), + } + ); + } + + #[test] + fn sandbox_owner_reference_accepts_v1alpha1_owner() { + let pod = pod_with_owner_refs(vec![sandbox_owner_with_api_version( + SANDBOX_API_VERSION_FULL_V1ALPHA1, + "sandbox-a", + "cr-uid-a", + )]); + + let owner = sandbox_owner_reference(&pod).expect("expected v1alpha1 Sandbox owner"); + + assert_eq!( + owner, + SandboxOwnerReference { + api_version: SANDBOX_API_VERSION_FULL_V1ALPHA1.to_string(), + name: "sandbox-a".to_string(), + uid: "cr-uid-a".to_string(), + } + ); + } + + #[test] + fn sandbox_owner_reference_rejects_missing_owner() { + let pod = pod_with_owner_refs(vec![]); + + let err = sandbox_owner_reference(&pod).expect_err("missing owner must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn sandbox_owner_reference_rejects_unsupported_sandbox_api_version() { + let pod = pod_with_owner_refs(vec![sandbox_owner_with_api_version( + "agents.x-k8s.io/v1", + "sandbox-a", + "cr-uid-a", + )]); + + let err = + sandbox_owner_reference(&pod).expect_err("unsupported apiVersion must fail closed"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn sandbox_owner_reference_requires_controlling_owner() { + let mut owner = sandbox_owner("sandbox-a", "cr-uid-a"); + owner.controller = Some(false); + let pod = pod_with_owner_refs(vec![owner]); + + let err = sandbox_owner_reference(&pod).expect_err("non-controller owner must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn sandbox_owner_reference_rejects_ambiguous_sandbox_owners() { + let pod = pod_with_owner_refs(vec![ + sandbox_owner("sandbox-a", "cr-uid-a"), + sandbox_owner("sandbox-b", "cr-uid-b"), + ]); + + let err = sandbox_owner_reference(&pod).expect_err("multiple owners must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn validate_sandbox_owner_reference_requires_matching_cr_uid() { + let owner = SandboxOwnerReference { + api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), + name: "sandbox-a".to_string(), + uid: "cr-uid-a".to_string(), + }; + let cr = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-a"); + validate_sandbox_owner_reference(&owner, &cr).expect("matching CR should be accepted"); + + let wrong_uid = sandbox_cr("sandbox-a", "cr-uid-b", "sandbox-id-a"); + let err = validate_sandbox_owner_reference(&owner, &wrong_uid) + .expect_err("wrong CR UID must fail"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[test] + fn validate_sandbox_owner_binding_requires_matching_label() { + let owner = SandboxOwnerReference { + api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), + name: "sandbox-a".to_string(), + uid: "cr-uid-a".to_string(), + }; + let cr = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-a"); + validate_sandbox_owner_binding(&owner, "sandbox-id-a", &cr) + .expect("matching label should be accepted"); + let wrong_label = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-b"); + let err = validate_sandbox_owner_binding(&owner, "sandbox-id-a", &wrong_label) + .expect_err("wrong sandbox-id label must fail"); + assert_eq!(err.code(), tonic::Code::PermissionDenied); + } + + #[tokio::test] + async fn provider_delegates_to_resolver() { + let resolved = bootstrap_identity(SupervisorBootstrapBinding::BoundSandbox { + sandbox_id: "sandbox-a".to_string(), + }); + let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved.clone())))); + let provider = KubernetesSupervisorBootstrapIdentityProvider::new(fake.clone()); + + let result = provider + .authenticate_registration("sa-jwt") + .await + .unwrap() + .expect("expected identity"); + + assert_eq!(result, resolved); + assert_eq!(fake.seen_tokens.lock().unwrap().as_slice(), ["sa-jwt"]); + } + + #[tokio::test] + async fn provider_allows_none_to_fall_through() { + let fake = Arc::new(FakeResolver::returning(Ok(None))); + let provider = KubernetesSupervisorBootstrapIdentityProvider::new(fake); + let result = provider.authenticate_registration("unknown").await.unwrap(); + assert!(result.is_none()); + } + + #[tokio::test] + async fn provider_accepts_warm_pending_identity() { + let resolved = bootstrap_identity(SupervisorBootstrapBinding::WarmPending); + let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved.clone())))); + let provider = KubernetesSupervisorBootstrapIdentityProvider::new(fake); + let identity = provider + .authenticate_registration("sa-jwt") + .await + .unwrap() + .expect("expected identity"); + + assert_eq!(identity.binding, SupervisorBootstrapBinding::WarmPending); + assert_eq!(identity.owner_uid, "cr-uid-a"); + } + + #[tokio::test] + async fn resolver_error_propagates() { + let fake = Arc::new(FakeResolver::returning(Err(Status::unavailable( + "apiserver down", + )))); + let provider = KubernetesSupervisorBootstrapIdentityProvider::new(fake); + let err = provider + .authenticate_registration("sa-jwt") + .await + .expect_err("resolver error must propagate"); + assert_eq!(err.code(), tonic::Code::Unavailable); + } +} diff --git a/crates/openshell-driver-kubernetes/src/config.rs b/crates/openshell-driver-kubernetes/src/config.rs index 1eeaac8396..32bccd940f 100644 --- a/crates/openshell-driver-kubernetes/src/config.rs +++ b/crates/openshell-driver-kubernetes/src/config.rs @@ -18,6 +18,12 @@ pub const DEFAULT_WORKSPACE_STORAGE_SIZE: &str = "2Gi"; /// Default non-root UID for relaxed Kubernetes network supervisor sidecars. pub const DEFAULT_PROXY_UID: u32 = 1337; +/// Label selector for ConfigMap-backed warm-pool profiles. +pub const DEFAULT_WARM_POOL_PROFILE_LABEL_SELECTOR: &str = "openshell.ai/warm-pool-profile=true"; + +/// Conservative upper bound for profile-declared warm-pool replicas. +pub const DEFAULT_WARM_POOL_PROFILE_MAX_REPLICAS: u32 = 100; + /// How the supervisor binary is delivered into sandbox pods. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] @@ -125,6 +131,82 @@ impl KubernetesSidecarConfig { } } +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct KubernetesWarmPoolingConfig { + /// Allow the Kubernetes driver to satisfy compatible create requests by + /// creating v1beta1 Agent Sandbox `SandboxClaim` resources. + pub enabled: bool, + /// Optional ConfigMap-backed warm-pool profile reconciler. + pub profiles: KubernetesWarmPoolProfilesConfig, +} + +impl Default for KubernetesWarmPoolingConfig { + fn default() -> Self { + Self { + enabled: true, + profiles: KubernetesWarmPoolProfilesConfig::default(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct KubernetesWarmPoolProfilesConfig { + /// Watch labelled `ConfigMap` resources and reconcile generated + /// `SandboxTemplate` and `SandboxWarmPool` resources from their TOML + /// profiles. + pub enabled: bool, + /// Namespace containing admin-authored warm-pool profile `ConfigMap` + /// resources. When empty, the driver uses its configured Kubernetes + /// namespace. + pub namespace: String, + /// Label selector used to select warm-pool profile `ConfigMap` resources. + pub label_selector: String, + /// Maximum allowed `replicas` value in a profile. + pub max_replicas: u32, +} + +impl Default for KubernetesWarmPoolProfilesConfig { + fn default() -> Self { + Self { + enabled: true, + namespace: String::new(), + label_selector: DEFAULT_WARM_POOL_PROFILE_LABEL_SELECTOR.to_string(), + max_replicas: DEFAULT_WARM_POOL_PROFILE_MAX_REPLICAS, + } + } +} + +impl KubernetesWarmPoolProfilesConfig { + #[must_use] + pub fn effective_namespace<'a>(&'a self, fallback_namespace: &'a str) -> &'a str { + if self.namespace.trim().is_empty() { + fallback_namespace + } else { + self.namespace.as_str() + } + } + + #[must_use] + pub fn effective_label_selector(&self) -> &str { + if self.label_selector.trim().is_empty() { + DEFAULT_WARM_POOL_PROFILE_LABEL_SELECTOR + } else { + self.label_selector.as_str() + } + } + + #[must_use] + pub fn effective_max_replicas(&self) -> u32 { + if self.max_replicas == 0 { + DEFAULT_WARM_POOL_PROFILE_MAX_REPLICAS + } else { + self.max_replicas + } + } +} + /// Kubernetes `AppArmor` profile requested for the sandbox agent container. #[derive(Debug, Clone, PartialEq, Eq)] pub enum AppArmorProfile { @@ -253,6 +335,8 @@ pub struct KubernetesComputeConfig { pub topology: SupervisorTopology, /// Sidecar-only settings used when `topology = "sidecar"`. pub sidecar: KubernetesSidecarConfig, + /// Warm-pool allocation settings. + pub warm_pooling: KubernetesWarmPoolingConfig, pub grpc_endpoint: String, pub ssh_socket_path: String, pub client_tls_secret_name: String, @@ -272,9 +356,9 @@ pub struct KubernetesComputeConfig { /// Empty string (default) = omit the field, using the cluster default. pub default_runtime_class_name: String, /// Lifetime (seconds) of the projected `ServiceAccount` token kubelet - /// writes into each sandbox pod. Used only for the one-shot - /// `IssueSandboxToken` bootstrap exchange — the gateway-minted JWT - /// that follows has its own TTL set via `gateway_jwt.ttl_secs`. + /// writes into each sandbox pod. Used only for the + /// `RegisterSupervisorPod` bootstrap stream — the gateway-minted JWT that + /// follows has its own TTL set via `gateway_jwt.ttl_secs`. /// /// Kubelet enforces a minimum of 600 seconds; the supervisor uses /// this token within a few seconds of pod start, so any value at @@ -340,6 +424,7 @@ impl Default for KubernetesComputeConfig { supervisor_sideload_method: SupervisorSideloadMethod::default(), topology: SupervisorTopology::default(), sidecar: KubernetesSidecarConfig::default(), + warm_pooling: KubernetesWarmPoolingConfig::default(), grpc_endpoint: String::new(), ssh_socket_path: "/run/openshell/ssh.sock".to_string(), client_tls_secret_name: String::new(), @@ -533,6 +618,57 @@ mod tests { assert!(cfg.sidecar.process_binary_aware_network_policy); } + #[test] + fn default_warm_pooling_is_enabled() { + let cfg = KubernetesComputeConfig::default(); + assert!(cfg.warm_pooling.enabled); + assert!(cfg.warm_pooling.profiles.enabled); + } + + #[test] + fn serde_override_warm_pooling_enabled() { + let json = serde_json::json!({ + "warm_pooling": { + "enabled": false + } + }); + let cfg: KubernetesComputeConfig = serde_json::from_value(json).unwrap(); + assert!(!cfg.warm_pooling.enabled); + } + + #[test] + fn serde_override_warm_pool_profile_reconciliation() { + let json = serde_json::json!({ + "warm_pooling": { + "profiles": { + "enabled": false, + "namespace": "profiles", + "label_selector": "openshell.ai/warm-pool-profile=true,tier=gpu", + "max_replicas": 7 + } + } + }); + let cfg: KubernetesComputeConfig = serde_json::from_value(json).unwrap(); + assert!(!cfg.warm_pooling.profiles.enabled); + assert_eq!(cfg.warm_pooling.profiles.namespace, "profiles"); + assert_eq!( + cfg.warm_pooling.profiles.label_selector, + "openshell.ai/warm-pool-profile=true,tier=gpu" + ); + assert_eq!(cfg.warm_pooling.profiles.max_replicas, 7); + } + + #[test] + fn serde_rejects_unknown_warm_pooling_field() { + let json = serde_json::json!({ + "warm_pooling": { + "mode": "always" + } + }); + let err = serde_json::from_value::(json).unwrap_err(); + assert!(err.to_string().contains("unknown field")); + } + #[test] fn serde_override_topology_sidecar() { let json = serde_json::json!({ diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index c784f10db9..8786a513e6 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -9,11 +9,15 @@ use crate::config::{ DEFAULT_WORKSPACE_STORAGE_SIZE, KubernetesComputeConfig, SupervisorSideloadMethod, SupervisorTopology, }; -use futures::{Stream, StreamExt, TryStreamExt}; +use crate::sandboxclaim::sandbox_claim_activation_token; +use futures::{Stream, StreamExt, TryStreamExt, stream}; use k8s_openapi::api::core::v1::{ - Event as KubeEventObj, Namespace, Node, PersistentVolumeClaimVolumeSource, Volume, VolumeMount, + ConfigMap, Event as KubeEventObj, Namespace, Node, PersistentVolumeClaimVolumeSource, Volume, + VolumeMount, +}; +use kube::api::{ + Api, ApiResource, DeleteParams, ListParams, Patch, PatchParams, PostParams, Preconditions, }; -use kube::api::{Api, ApiResource, DeleteParams, ListParams, PostParams, Preconditions}; use kube::core::gvk::GroupVersionKind; use kube::core::{DynamicObject, ObjectMeta}; use kube::runtime::watcher::{self, Event}; @@ -30,20 +34,21 @@ use openshell_core::progress::{ }; use openshell_core::proto::compute::v1::{ DriverCondition as SandboxCondition, DriverPlatformEvent as PlatformEvent, - DriverSandbox as Sandbox, DriverSandboxSpec as SandboxSpec, + DriverResourceRequirements, DriverSandbox as Sandbox, DriverSandboxSpec as SandboxSpec, DriverSandboxStatus as SandboxStatus, DriverSandboxTemplate as SandboxTemplate, - GetCapabilitiesResponse, GpuResourceRequirements, WatchSandboxesDeletedEvent, - WatchSandboxesEvent, WatchSandboxesPlatformEvent, WatchSandboxesSandboxEvent, - watch_sandboxes_event, + GetCapabilitiesResponse, GpuResourceRequirements, ResourceRequirements, + WatchSandboxesDeletedEvent, WatchSandboxesEvent, WatchSandboxesPlatformEvent, + WatchSandboxesSandboxEvent, watch_sandboxes_event, }; use openshell_core::proto_struct::{struct_to_json_object, value_to_json}; use serde::Deserialize; +use sha2::{Digest, Sha256}; use std::collections::{BTreeMap, HashSet}; use std::path::Path; use std::pin::Pin; use std::sync::Arc; use std::time::Duration; -use tokio::sync::{OnceCell, mpsc}; +use tokio::sync::{OnceCell, RwLock, mpsc}; use tokio_stream::wrappers::ReceiverStream; use tracing::{debug, info, warn}; @@ -86,6 +91,7 @@ impl From for openshell_core::ComputeDriverError { /// This prevents gRPC handlers from blocking indefinitely when the k8s /// API server is unreachable or slow. const KUBE_API_TIMEOUT: Duration = Duration::from_secs(30); +const WARM_POOL_CACHE_RETRY_DELAY: Duration = Duration::from_secs(10); const SANDBOX_GROUP: &str = "agents.x-k8s.io"; const SANDBOX_VERSION_V1BETA1: &str = "v1beta1"; @@ -93,6 +99,37 @@ const SANDBOX_VERSION_V1ALPHA1: &str = "v1alpha1"; const SANDBOX_VERSIONS: &[&str] = &[SANDBOX_VERSION_V1BETA1, SANDBOX_VERSION_V1ALPHA1]; pub const SANDBOX_KIND: &str = "Sandbox"; +const EXTENSIONS_GROUP: &str = "extensions.agents.x-k8s.io"; +const EXTENSIONS_VERSION_V1BETA1: &str = "v1beta1"; +const SANDBOX_CLAIM_KIND: &str = "SandboxClaim"; +const SANDBOX_WARM_POOL_KIND: &str = "SandboxWarmPool"; +const SANDBOX_TEMPLATE_KIND: &str = "SandboxTemplate"; +const LABEL_WARM_POOL_ENABLED: &str = "openshell.ai/enabled"; +const LABEL_ALLOCATION: &str = "openshell.ai/allocation"; +const LABEL_ALLOCATION_SANDBOX_CLAIM: &str = "sandbox-claim"; +const LABEL_WARM_POOL_PROFILE: &str = "openshell.ai/warm-pool-profile"; +const LABEL_WARM_POOL_PROFILE_UID: &str = "openshell.ai/warm-pool-profile-uid"; +const LABEL_WARM_POOL_MANAGED_BY: &str = "openshell.ai/managed-by"; +const LABEL_WARM_POOL_MANAGED_BY_VALUE: &str = "openshell-kubernetes-driver"; +const ANNOTATION_WARM_POOL_PROFILE_NAMESPACE: &str = "openshell.ai/warm-pool-profile-namespace"; +const ANNOTATION_WARM_POOL_PROFILE_NAME: &str = "openshell.ai/warm-pool-profile-name"; +const ANNOTATION_WARM_POOL_PROFILE_UID: &str = "openshell.ai/warm-pool-profile-uid"; +const ANNOTATION_WARM_POOL_SOURCE_RESOURCE_VERSION: &str = "openshell.ai/source-resource-version"; +const ANNOTATION_WARM_POOL_TEMPLATE_FINGERPRINT: &str = "openshell.ai/template-fingerprint"; +const ANNOTATION_WARM_POOL_STATUS: &str = "openshell.ai/warm-pool-status"; +const ANNOTATION_WARM_POOL_OBSERVED_RESOURCE_VERSION: &str = + "openshell.ai/warm-pool-observed-resource-version"; +const ANNOTATION_WARM_POOL_TARGET_NAMESPACE: &str = "openshell.ai/warm-pool-target-namespace"; +const ANNOTATION_WARM_POOL_TEMPLATE: &str = "openshell.ai/warm-pool-template"; +const ANNOTATION_WARM_POOL_FINGERPRINT: &str = "openshell.ai/warm-pool-fingerprint"; +const ANNOTATION_WARM_POOL_LAST_ERROR: &str = "openshell.ai/warm-pool-last-error"; +const POD_ANNOTATION_SANDBOX_ID: &str = "openshell.io/sandbox-id"; +const WARM_POOL_PROFILE_DATA_KEY: &str = "warm-pool.toml"; +const WARM_POOL_PROFILE_VERSION: u32 = 1; +const WARM_POOL_PROFILE_REQUEUE_DELAY: Duration = Duration::from_secs(10); +const WARM_POOL_PROFILE_RESYNC_DELAY: Duration = Duration::from_secs(300); +const WARM_POOL_PROFILE_NAME_PREFIX: &str = "openshell-wp"; + const GPU_RESOURCE_NAME: &str = "nvidia.com/gpu"; const SPIFFE_WORKLOAD_API_VOLUME_NAME: &str = "spiffe-workload-api"; @@ -101,6 +138,150 @@ struct AgentSandboxApi { resource: ApiResource, } +struct ExtensionApi { + api: Api, + resource: ApiResource, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +struct EnabledWarmPool { + namespace: String, + name: String, + template_namespace: String, + template_name: String, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +struct WarmPoolCacheEntry { + namespace: String, + name: String, + template_namespace: String, + template_name: String, + template_fingerprint: String, +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct WarmPoolProfile { + version: u32, + workspace: String, + #[serde(default = "default_warm_pool_profile_replicas")] + replicas: u32, + #[serde(default)] + image: String, + #[serde(default)] + runtime_class_name: String, + #[serde(default)] + environment: BTreeMap, + #[serde(default)] + resources: WarmPoolProfileResources, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +struct WarmPoolProfileResources { + #[serde(default)] + cpu: String, + #[serde(default)] + memory: String, + gpu_count: Option, +} + +fn default_warm_pool_profile_replicas() -> u32 { + 1 +} + +#[derive(Debug, Clone)] +struct WarmPoolProfileSource { + namespace: String, + name: String, + uid: String, + resource_version: String, +} + +#[derive(Debug, Clone)] +struct RenderedWarmPoolProfile { + source: WarmPoolProfileSource, + workspace: String, + target_namespace: String, + generated_name: String, + replicas: u32, + template_spec: serde_json::Value, + fingerprint: String, +} + +#[derive(Debug, Clone, Default)] +struct WarmPoolCache { + state: Arc>, +} + +#[derive(Debug, Default)] +struct WarmPoolCacheState { + ready: bool, + entries: BTreeMap<(String, String), WarmPoolCacheEntry>, +} + +#[derive(Debug, Eq, PartialEq)] +enum WarmPoolCacheLookup { + NotReady, + NoMatch, + Match(WarmPoolCacheEntry), + Ambiguous(usize), +} + +impl WarmPoolCache { + async fn replace_entries(&self, entries: Vec) { + let mut state = self.state.write().await; + state.ready = true; + state.entries = entries + .into_iter() + .map(|entry| ((entry.namespace.clone(), entry.name.clone()), entry)) + .collect(); + } + + async fn mark_not_ready(&self) { + let mut state = self.state.write().await; + state.ready = false; + state.entries.clear(); + } + + async fn matching_pool( + &self, + namespace: &str, + template_fingerprint: &str, + ) -> WarmPoolCacheLookup { + let state = self.state.read().await; + if !state.ready { + return WarmPoolCacheLookup::NotReady; + } + + let mut matches = state + .entries + .values() + .filter(|entry| { + entry.namespace == namespace && entry.template_fingerprint == template_fingerprint + }) + .cloned() + .collect::>(); + + match matches.len() { + 0 => WarmPoolCacheLookup::NoMatch, + 1 => WarmPoolCacheLookup::Match(matches.pop().expect("one match exists")), + count => WarmPoolCacheLookup::Ambiguous(count), + } + } + + async fn entries_for_namespace(&self, namespace: &str) -> Vec { + let state = self.state.read().await; + state + .entries + .values() + .filter(|entry| entry.namespace == namespace) + .cloned() + .collect() + } +} + // This POC treats the selected Struct as a driver-local typed schema. Once the // Kubernetes shape stabilizes, these serde structs may move to driver-local // protobuf definitions, but the typed decode should stay inside this driver. @@ -431,6 +612,7 @@ pub struct KubernetesComputeDriver { client: Client, watch_client: Client, sandbox_api_version: Arc>, + warm_pool_cache: WarmPoolCache, config: KubernetesComputeConfig, } @@ -476,12 +658,20 @@ impl KubernetesComputeDriver { let watch_client = Client::try_from(watch_kube_config).map_err(KubernetesDriverError::from_kube)?; - Ok(Self { + let driver = Self { client, watch_client, sandbox_api_version: Arc::new(OnceCell::new()), + warm_pool_cache: WarmPoolCache::default(), config, - }) + }; + if driver.config.warm_pooling.enabled { + driver.spawn_warm_pool_cache_controller(); + if driver.config.warm_pooling.profiles.enabled { + driver.spawn_warm_pool_profile_reconciler(); + } + } + Ok(driver) } pub fn capabilities(&self) -> Result { @@ -525,6 +715,64 @@ impl KubernetesComputeDriver { AgentSandboxApi { api, resource } } + fn extension_resource(kind: &str) -> ApiResource { + let gvk = GroupVersionKind::gvk(EXTENSIONS_GROUP, EXTENSIONS_VERSION_V1BETA1, kind); + ApiResource::from_gvk(&gvk) + } + + fn namespaced_extension_api(client: Client, namespace: &str, kind: &str) -> ExtensionApi { + let resource = Self::extension_resource(kind); + let api = Api::namespaced_with(client, namespace, &resource); + ExtensionApi { api, resource } + } + + fn all_extension_api(client: Client, kind: &str) -> ExtensionApi { + let resource = Self::extension_resource(kind); + let api = Api::all_with(client, &resource); + ExtensionApi { api, resource } + } + + fn sandbox_claim_api(client: Client, namespace: &str) -> ExtensionApi { + Self::namespaced_extension_api(client, namespace, SANDBOX_CLAIM_KIND) + } + + fn sandbox_claim_api_all(client: Client) -> ExtensionApi { + Self::all_extension_api(client, SANDBOX_CLAIM_KIND) + } + + fn sandbox_warm_pool_api_all(client: Client) -> ExtensionApi { + Self::all_extension_api(client, SANDBOX_WARM_POOL_KIND) + } + + fn sandbox_warm_pool_api(client: Client, namespace: &str) -> ExtensionApi { + Self::namespaced_extension_api(client, namespace, SANDBOX_WARM_POOL_KIND) + } + + fn sandbox_template_api(client: Client, namespace: &str) -> ExtensionApi { + Self::namespaced_extension_api(client, namespace, SANDBOX_TEMPLATE_KIND) + } + + fn spawn_warm_pool_cache_controller(&self) { + let cache = self.warm_pool_cache.clone(); + let client = self.client.clone(); + let watch_client = self.watch_client.clone(); + let fallback_namespace = self.config.namespace.clone(); + + tokio::spawn(async move { + run_warm_pool_cache_controller(cache, client, watch_client, fallback_namespace).await; + }); + } + + fn spawn_warm_pool_profile_reconciler(&self) { + let client = self.client.clone(); + let watch_client = self.watch_client.clone(); + let config = self.config.clone(); + + tokio::spawn(async move { + run_warm_pool_profile_reconciler(client, watch_client, config).await; + }); + } + async fn supported_agent_sandbox_api(&self, client: Client) -> Result { let sandbox_api_version = self.supported_sandbox_api_version(client.clone()).await?; Ok(self.agent_sandbox_api(client, sandbox_api_version)) @@ -591,66 +839,12 @@ impl KubernetesComputeDriver { /// annotations. /// - If neither config nor `OpenShift` is found, returns `(1000, 1000, {})` as defaults. async fn resolve_sandbox_identity(&self) -> (u32, u32, BTreeMap) { - // Explicit config takes priority — skip namespace lookup entirely. - if self.config.sandbox_uid.is_some() { - let uid = self.config.resolve_sandbox_uid(None); - let gid = self.config.resolve_sandbox_gid(uid, None); - return (uid, gid, BTreeMap::new()); - } - - // Try to read namespace annotations for OpenShift SCC. - // Namespace is namespaced so Api::all works (it's cluster-scoped but - // can list all namespaces) and we filter by name, or use Api::namespaced. - let ns_api: Api = Api::all(self.client.clone()); - match tokio::time::timeout(KUBE_API_TIMEOUT, ns_api.get(self.config.namespace.as_str())) - .await - { - Ok(Ok(ns)) => { - let anns = ns.metadata.annotations.unwrap_or_default(); - tracing::info!( - namespace = %self.config.namespace, - uid_range = ?anns.get(crate::config::ANNOTATION_SCC_UID_RANGE), - sup_groups = ?anns.get(crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS), - "Resolved namespace annotations for sandbox identity" - ); - let uid = self.config.resolve_sandbox_uid(Some(&anns)); - // Explicit sandbox_gid config wins; SCC annotation only applies when not set. - let baseline_gid = self.config.resolve_sandbox_gid(uid, None); - let gid = self.config.sandbox_gid.map_or_else( - || { - anns.get(crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS) - .and_then(|sup_range| { - KubernetesComputeConfig::from_open_shift_supplemental_groups( - sup_range, - ) - }) - .unwrap_or(baseline_gid) - }, - |_| baseline_gid, - ); - tracing::info!(uid, gid, "Resolved sandbox identity"); - (uid, gid, anns) - } - Ok(Err(e)) => { - tracing::warn!( - namespace = %self.config.namespace, - error = %e, - "Failed to fetch namespace for SCC annotations, falling back to defaults" - ); - let uid = DEFAULT_SANDBOX_UID; - let gid = self.config.resolve_sandbox_gid(uid, None); - (uid, gid, BTreeMap::new()) - } - Err(_) => { - tracing::warn!( - namespace = %self.config.namespace, - "Namespace fetch timed out, falling back to defaults" - ); - let uid = DEFAULT_SANDBOX_UID; - let gid = self.config.resolve_sandbox_gid(uid, None); - (uid, gid, BTreeMap::new()) - } - } + resolve_sandbox_identity_for_config( + self.client.clone(), + &self.config, + self.config.namespace.as_str(), + ) + .await } async fn has_gpu_capacity(&self) -> Result { @@ -700,17 +894,17 @@ impl KubernetesComputeDriver { format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE},{LABEL_SANDBOX_ID}={sandbox_id}"); let lp = ListParams::default().labels(&selector); match tokio::time::timeout(KUBE_API_TIMEOUT, agent_sandbox_api.api.list(&lp)).await { - Ok(Ok(list)) => list.items.into_iter().next().map_or_else( - || { - debug!(sandbox_id = %sandbox_id, "Sandbox not found in Kubernetes"); - Ok(None) - }, - |obj| { - Ok(sandbox_from_object(&self.config.namespace, obj) + Ok(Ok(list)) => { + if let Some(obj) = list.items.into_iter().next() { + return Ok(sandbox_from_object(&self.config.namespace, obj) .ok() - .map(|(_, s)| s)) - }, - ), + .map(|(_, s)| s)); + } + { + debug!(sandbox_id = %sandbox_id, "Sandbox not found in Kubernetes"); + self.get_claim_backed_sandbox(sandbox_id).await + } + } Ok(Err(err)) => { warn!( sandbox_id = %sandbox_id, @@ -765,6 +959,13 @@ impl KubernetesComputeDriver { } }) .collect(); + let mut claim_sandboxes = self.list_claim_backed_sandboxes().await?; + let mut claim_ids = claim_sandboxes + .iter() + .map(|sandbox| sandbox.id.clone()) + .collect::>(); + sandboxes.retain(|sandbox| !claim_ids.remove(&sandbox.id)); + sandboxes.append(&mut claim_sandboxes); sandboxes.sort_by(|left, right| { left.name .cmp(&right.name) @@ -794,115 +995,541 @@ impl KubernetesComputeDriver { } } - #[allow(clippy::similar_names)] - pub async fn create_sandbox(&self, sandbox: &Sandbox) -> Result<(), KubernetesDriverError> { - let gpu_requirements = sandbox - .spec - .as_ref() - .and_then(|spec| driver_gpu_requirements(spec.resource_requirements.as_ref())); - validate_gpu_request(gpu_requirements).map_err(|status| { - KubernetesDriverError::InvalidArgument(status.message().to_string()) - })?; - - // Validate sandbox name against Kubernetes naming requirements - validate_kubernetes_dns1123_label(&sandbox.name, "sandbox name") - .map_err(KubernetesDriverError::InvalidArgument)?; + async fn list_claim_backed_sandboxes(&self) -> Result, String> { + let selector = format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE}"); + let lp = ListParams::default().labels(&selector); + let claim_api = Self::sandbox_claim_api_all(self.client.clone()); + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(Vec::new()), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let claim_api = + Self::sandbox_claim_api(self.client.clone(), &self.config.namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(Vec::new()), + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; - let name = sandbox.name.as_str(); - info!( - sandbox_id = %sandbox.id, - sandbox_name = %name, - namespace = %self.config.namespace, - "Creating sandbox in Kubernetes" - ); + Ok(list + .items + .into_iter() + .filter_map(|obj| sandbox_from_claim_object(&self.config.namespace, obj).ok()) + .collect()) + } - let agent_sandbox_api = self - .supported_agent_sandbox_api(self.client.clone()) - .await - .map_err(KubernetesDriverError::Message)?; + async fn get_claim_backed_sandbox(&self, sandbox_id: &str) -> Result, String> { + let selector = + format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE},{LABEL_SANDBOX_ID}={sandbox_id}"); + let lp = ListParams::default().labels(&selector); + let claim_api = Self::sandbox_claim_api_all(self.client.clone()); + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(None), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let claim_api = + Self::sandbox_claim_api(self.client.clone(), &self.config.namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(None), + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; - // Resolve sandbox UID/GID from config or OpenShift SCC namespace annotations. - let (resolved_user_id, resolved_group_id, ns_annotations) = - self.resolve_sandbox_identity().await; + Ok(list + .items + .into_iter() + .find_map(|obj| sandbox_from_claim_object(&self.config.namespace, obj).ok())) + } - let params = SandboxPodParams { - default_image: &self.config.default_image, - image_pull_policy: &self.config.image_pull_policy, - image_pull_secrets: &self.config.image_pull_secrets, - supervisor_image: &self.config.supervisor_image, - supervisor_image_pull_policy: &self.config.supervisor_image_pull_policy, - supervisor_sideload_method: self.config.supervisor_sideload_method, - topology: self.config.topology, - proxy_uid: self.config.sidecar.proxy_uid, - process_binary_aware_network_policy: self - .config - .sidecar - .process_binary_aware_network_policy, - service_account_name: &self.config.service_account_name, - sandbox_id: &sandbox.id, - sandbox_name: &sandbox.name, - grpc_endpoint: &self.config.grpc_endpoint, - ssh_socket_path: self.ssh_socket_path(), - client_tls_secret_name: &self.config.client_tls_secret_name, - host_gateway_ip: &self.config.host_gateway_ip, - enable_user_namespaces: self.config.enable_user_namespaces, - app_armor_profile: self.config.app_armor_profile.as_ref(), - workspace_default_storage_size: &self.config.workspace_default_storage_size, - default_runtime_class_name: &self.config.default_runtime_class_name, - sa_token_ttl_secs: self.config.effective_sa_token_ttl_secs(), - provider_spiffe_enabled: self.config.provider_spiffe_enabled(), - provider_spiffe_workload_api_socket_path: &self - .config - .provider_spiffe_workload_api_socket_path, - sandbox_uid: resolved_user_id, - sandbox_gid: resolved_group_id, + async fn list_enabled_warm_pools_with_client( + client: Client, + fallback_namespace: &str, + ) -> Result, String> { + let lp = ListParams::default().labels(&format!("{LABEL_WARM_POOL_ENABLED}=true")); + let warm_pool_api = Self::sandbox_warm_pool_api_all(client.clone()); + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, warm_pool_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(Vec::new()), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let warm_pool_api = Self::sandbox_warm_pool_api(client, fallback_namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, warm_pool_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(Vec::new()), + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } }; - validate_sidecar_proxy_identity(¶ms)?; - let data = sandbox_to_k8s_spec(sandbox.spec.as_ref(), ¶ms) - .map_err(KubernetesDriverError::InvalidArgument)?; - let kube_name = kube_resource_name(&sandbox.workspace, name); - let mut obj = DynamicObject::new(&kube_name, &agent_sandbox_api.resource); - // Copy only the SCC-related annotations onto the Sandbox CR for - // traceability. Copying the full namespace annotation map exposes - // unrelated cluster metadata and can fail with oversized annotations. - let mut annotations = sandbox_annotations(sandbox); - for key in [ - crate::config::ANNOTATION_SCC_UID_RANGE, - crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS, - ] { - if let Some(v) = ns_annotations.get(key) { - annotations.insert(key.to_string(), v.clone()); - } + Ok(list + .items + .into_iter() + .filter_map(enabled_warm_pool_from_object) + .collect()) + } + + async fn template_fingerprint_for_warm_pool_with_client( + client: Client, + pool: &EnabledWarmPool, + ) -> Result, String> { + let template_api = Self::sandbox_template_api(client, &pool.template_namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, template_api.api.get(&pool.template_name)) + .await + { + Ok(Ok(template)) => sandbox_template_fingerprint(&template).map(Some), + Ok(Err(KubeError::Api(err))) if err.code == 404 => Ok(None), + Ok(Err(err)) if extension_api_unavailable(&err) => Ok(None), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), } - obj.metadata = ObjectMeta { - name: Some(kube_name), - namespace: Some(self.config.namespace.clone()), - labels: Some(sandbox_labels(sandbox)), - annotations: Some(annotations), - ..Default::default() - }; + } - obj.data = data; + async fn extension_api_for_watch( + client: Client, + fallback_namespace: &str, + kind: &str, + ) -> Result, String> { + let all_api = Self::all_extension_api(client.clone(), kind); match tokio::time::timeout( KUBE_API_TIMEOUT, - agent_sandbox_api.api.create(&PostParams::default(), &obj), + all_api.api.list(&ListParams::default().limit(1)), ) .await { - Ok(Ok(_result)) => { - info!( - sandbox_id = %sandbox.id, - sandbox_name = %name, - "Sandbox created in Kubernetes successfully" - ); - Ok(()) + Ok(Ok(_)) => Ok(all_api.api), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let namespaced_api = + Self::namespaced_extension_api(client, fallback_namespace, kind); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + namespaced_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => Ok(namespaced_api.api), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } } - Ok(Err(err)) => { - warn!( - sandbox_id = %sandbox.id, - sandbox_name = %name, - error = %err, + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } + } + + async fn matching_warm_pool( + &self, + target_namespace: &str, + request_fingerprint: &str, + ) -> Option { + match self + .warm_pool_cache + .matching_pool(target_namespace, request_fingerprint) + .await + { + WarmPoolCacheLookup::NotReady => { + debug!( + namespace = %target_namespace, + "Warm-pool cache is not ready; falling back to direct Sandbox" + ); + None + } + WarmPoolCacheLookup::NoMatch => { + let cached_warm_pool_templates = self + .warm_pool_cache + .entries_for_namespace(target_namespace) + .await + .into_iter() + .map(|entry| { + format!( + "{}/{}:{}", + entry.template_namespace, + entry.template_name, + entry.template_fingerprint + ) + }) + .collect::>(); + debug!( + namespace = %target_namespace, + request_fingerprint, + cached_warm_pool_template_count = cached_warm_pool_templates.len(), + cached_warm_pool_templates = ?cached_warm_pool_templates, + "No OpenShell-enabled warm pool matches sandbox request; falling back to direct Sandbox" + ); + None + } + WarmPoolCacheLookup::Match(warm_pool) => { + debug!( + namespace = %target_namespace, + request_fingerprint, + warm_pool = %warm_pool.name, + template_namespace = %warm_pool.template_namespace, + template = %warm_pool.template_name, + "OpenShell-enabled warm pool matches sandbox request" + ); + Some(warm_pool) + } + WarmPoolCacheLookup::Ambiguous(count) => { + warn!( + namespace = %target_namespace, + count, + "Multiple OpenShell-enabled warm pools match sandbox request; falling back to direct Sandbox" + ); + None + } + } + } + + async fn try_create_sandbox_claim( + &self, + sandbox: &Sandbox, + target_namespace: &str, + request_fingerprint: &str, + ) -> Result, KubernetesDriverError> { + if !self.config.warm_pooling.enabled { + return Ok(None); + } + + let Some(warm_pool) = self + .matching_warm_pool(target_namespace, request_fingerprint) + .await + else { + return Ok(None); + }; + + let claim_api = Self::sandbox_claim_api(self.client.clone(), target_namespace); + let claim = sandbox_claim_to_k8s_object(sandbox, &warm_pool.name, &claim_api.resource); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + claim_api.api.create(&PostParams::default(), &claim), + ) + .await + { + Ok(Ok(result)) => { + let claim_name = result + .metadata + .name + .as_deref() + .filter(|name| !name.is_empty()) + .unwrap_or(sandbox.name.as_str()); + let claim_uid = result + .metadata + .uid + .as_deref() + .filter(|uid| !uid.is_empty()) + .ok_or_else(|| { + KubernetesDriverError::Precondition( + "created SandboxClaim is missing Kubernetes UID".to_string(), + ) + })?; + let activation_token = + sandbox_claim_activation_token(target_namespace, claim_name, claim_uid); + info!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + warm_pool = %warm_pool.name, + "SandboxClaim created in Kubernetes successfully" + ); + Ok(Some(activation_token)) + } + Ok(Err(err)) if extension_api_unavailable(&err) => { + debug!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + error = %err, + "SandboxClaim API is unavailable; falling back to direct Sandbox" + ); + Ok(None) + } + Ok(Err(err)) if matches!(&err, KubeError::Api(api) if api.code == 409) => { + Err(KubernetesDriverError::from_kube(err)) + } + Ok(Err(err)) => { + warn!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + warm_pool = %warm_pool.name, + error = %err, + "Failed to create SandboxClaim; falling back to direct Sandbox" + ); + Ok(None) + } + Err(_elapsed) => { + warn!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + timeout_secs = KUBE_API_TIMEOUT.as_secs(), + "Timed out creating SandboxClaim; falling back to direct Sandbox" + ); + Ok(None) + } + } + } + + async fn delete_sandbox_claim(&self, sandbox_id: &str) -> Result { + let selector = + format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE},{LABEL_SANDBOX_ID}={sandbox_id}"); + let lp = ListParams::default().labels(&selector); + let mut claim_api = Self::sandbox_claim_api_all(self.client.clone()); + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(false), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + claim_api = Self::sandbox_claim_api(self.client.clone(), &self.config.namespace); + match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.list(&lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) if extension_api_unavailable(&err) => return Ok(false), + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + + let Some(obj) = list.items.into_iter().next() else { + return Ok(false); + }; + let Some(claim_name) = obj.metadata.name.clone() else { + return Ok(false); + }; + let claim_namespace = obj + .metadata + .namespace + .clone() + .unwrap_or_else(|| self.config.namespace.clone()); + let preconditions = Preconditions { + uid: obj.metadata.uid, + resource_version: obj.metadata.resource_version, + }; + + let claim_api = Self::sandbox_claim_api(self.client.clone(), &claim_namespace); + let dp = DeleteParams::default().preconditions(preconditions); + match tokio::time::timeout(KUBE_API_TIMEOUT, claim_api.api.delete(&claim_name, &dp)).await { + Ok(Ok(_response)) => { + info!( + sandbox_id = %sandbox_id, + namespace = %claim_namespace, + sandbox_claim = %claim_name, + "SandboxClaim deleted from Kubernetes" + ); + Ok(true) + } + Ok(Err(KubeError::Api(err))) if err.code == 404 || err.code == 409 => Ok(false), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } + } + + #[allow(clippy::similar_names)] + pub async fn create_sandbox( + &self, + sandbox: &Sandbox, + ) -> Result, KubernetesDriverError> { + let gpu_requirements = sandbox + .spec + .as_ref() + .and_then(|spec| driver_gpu_requirements(spec.resource_requirements.as_ref())); + validate_gpu_request(gpu_requirements).map_err(|status| { + KubernetesDriverError::InvalidArgument(status.message().to_string()) + })?; + + // Validate sandbox name against Kubernetes naming requirements + validate_kubernetes_dns1123_label(&sandbox.name, "sandbox name") + .map_err(KubernetesDriverError::InvalidArgument)?; + + let name = sandbox.name.as_str(); + info!( + sandbox_id = %sandbox.id, + sandbox_name = %name, + namespace = %self.config.namespace, + "Creating sandbox in Kubernetes" + ); + + let agent_sandbox_api = self + .supported_agent_sandbox_api(self.client.clone()) + .await + .map_err(KubernetesDriverError::Message)?; + + // Resolve sandbox UID/GID from config or OpenShift SCC namespace annotations. + let (resolved_user_id, resolved_group_id, ns_annotations) = + self.resolve_sandbox_identity().await; + + let params = SandboxPodParams { + default_image: &self.config.default_image, + image_pull_policy: &self.config.image_pull_policy, + image_pull_secrets: &self.config.image_pull_secrets, + supervisor_image: &self.config.supervisor_image, + supervisor_image_pull_policy: &self.config.supervisor_image_pull_policy, + supervisor_sideload_method: self.config.supervisor_sideload_method, + topology: self.config.topology, + proxy_uid: self.config.sidecar.proxy_uid, + process_binary_aware_network_policy: self + .config + .sidecar + .process_binary_aware_network_policy, + service_account_name: &self.config.service_account_name, + sandbox_id: &sandbox.id, + sandbox_name: &sandbox.name, + grpc_endpoint: &self.config.grpc_endpoint, + ssh_socket_path: self.ssh_socket_path(), + client_tls_secret_name: &self.config.client_tls_secret_name, + host_gateway_ip: &self.config.host_gateway_ip, + enable_user_namespaces: self.config.enable_user_namespaces, + app_armor_profile: self.config.app_armor_profile.as_ref(), + workspace_default_storage_size: &self.config.workspace_default_storage_size, + default_runtime_class_name: &self.config.default_runtime_class_name, + sa_token_ttl_secs: self.config.effective_sa_token_ttl_secs(), + provider_spiffe_enabled: self.config.provider_spiffe_enabled(), + provider_spiffe_workload_api_socket_path: &self + .config + .provider_spiffe_workload_api_socket_path, + sandbox_uid: resolved_user_id, + sandbox_gid: resolved_group_id, + }; + validate_sidecar_proxy_identity(¶ms)?; + + let data = sandbox_to_k8s_spec(sandbox.spec.as_ref(), ¶ms) + .map_err(KubernetesDriverError::InvalidArgument)?; + let target_namespace = self.config.namespace.clone(); + if self.config.warm_pooling.enabled { + match sandbox_spec_fingerprint(&data) { + Ok(fingerprint) => { + if let Some(activation_token) = self + .try_create_sandbox_claim(sandbox, &target_namespace, &fingerprint) + .await? + { + return Ok(Some(activation_token)); + } + debug!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + request_fingerprint = %fingerprint, + normalized_request_spec = %sandbox_spec_fingerprint_debug_json(&data) + .unwrap_or_else(|err| format!("failed to render normalized request spec: {err}")), + "Warm-pool request fingerprint input" + ); + } + Err(err) => { + debug!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + error = %err, + "Could not fingerprint sandbox request for warm-pool matching; falling back to direct Sandbox" + ); + } + } + } + let kube_name = kube_resource_name(&sandbox.workspace, name); + let mut obj = DynamicObject::new(&kube_name, &agent_sandbox_api.resource); + // Copy only the SCC-related annotations onto the Sandbox CR for + // traceability. Copying the full namespace annotation map exposes + // unrelated cluster metadata and can fail with oversized annotations. + let mut annotations = sandbox_annotations(sandbox); + for key in [ + crate::config::ANNOTATION_SCC_UID_RANGE, + crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS, + ] { + if let Some(v) = ns_annotations.get(key) { + annotations.insert(key.to_string(), v.clone()); + } + } + obj.metadata = ObjectMeta { + name: Some(kube_name), + namespace: Some(target_namespace), + labels: Some(sandbox_labels(sandbox)), + annotations: Some(annotations), + ..Default::default() + }; + + obj.data = data; + match tokio::time::timeout( + KUBE_API_TIMEOUT, + agent_sandbox_api.api.create(&PostParams::default(), &obj), + ) + .await + { + Ok(Ok(_result)) => { + info!( + sandbox_id = %sandbox.id, + sandbox_name = %name, + "Sandbox created in Kubernetes successfully" + ); + Ok(None) + } + Ok(Err(err)) => { + warn!( + sandbox_id = %sandbox.id, + sandbox_name = %name, + error = %err, "Failed to create sandbox in Kubernetes" ); Err(KubernetesDriverError::from_kube(err)) @@ -929,6 +1556,10 @@ impl KubernetesComputeDriver { "Deleting sandbox from Kubernetes" ); + if self.delete_sandbox_claim(sandbox_id).await? { + return Ok(true); + } + let agent_sandbox_api = self .supported_agent_sandbox_api(self.client.clone()) .await?; @@ -1024,7 +1655,12 @@ impl KubernetesComputeDriver { format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE},{LABEL_SANDBOX_ID}={sandbox_id}"); let lp = ListParams::default().labels(&selector); match tokio::time::timeout(KUBE_API_TIMEOUT, agent_sandbox_api.api.list(&lp)).await { - Ok(Ok(list)) => Ok(!list.items.is_empty()), + Ok(Ok(list)) => { + if !list.items.is_empty() { + return Ok(true); + } + Ok(self.get_claim_backed_sandbox(sandbox_id).await?.is_some()) + } Ok(Err(err)) => Err(err.to_string()), Err(_elapsed) => Err(format!( "timed out after {}s waiting for Kubernetes API", @@ -1043,6 +1679,53 @@ impl KubernetesComputeDriver { let event_api: Api = Api::namespaced(self.watch_client.clone(), &namespace); let watcher_config = watcher::Config::default().labels(&openshell_sandbox_label_selector()); let mut sandbox_stream = watcher::watcher(agent_sandbox_api.api, watcher_config).boxed(); + let claim_watcher_config = watcher::Config::default() + .labels(&format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE}")); + let claim_api = if self.config.warm_pooling.enabled { + let all_api = Self::sandbox_claim_api_all(self.watch_client.clone()); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + all_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => Some(all_api.api), + Ok(Err(KubeError::Api(err))) if err.code == 403 => { + let namespaced_api = + Self::sandbox_claim_api(self.watch_client.clone(), &namespace); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + namespaced_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => Some(namespaced_api.api), + Ok(Err(err)) if extension_api_unavailable(&err) => None, + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(Err(err)) if extension_api_unavailable(&err) => None, + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } else { + None + }; + let mut claim_stream = claim_api.map_or_else( + || stream::pending().boxed(), + |api| watcher::watcher(api, claim_watcher_config).boxed(), + ); let mut event_stream = watcher::watcher(event_api, watcher::Config::default()).boxed(); let (tx, rx) = mpsc::channel(256); @@ -1107,16 +1790,12 @@ impl KubernetesComputeDriver { break; } }, - result = event_stream.try_next() => match result { + result = claim_stream.try_next() => match result { Ok(Some(Event::Applied(obj))) => { - if let Some((sandbox_id, event)) = map_kube_event_to_platform( - &sandbox_name_to_id, - &agent_pod_to_id, - &obj, - ) { + if let Ok(sandbox) = sandbox_from_claim_object(&namespace, obj) { let event = WatchSandboxesEvent { - payload: Some(watch_sandboxes_event::Payload::PlatformEvent( - WatchSandboxesPlatformEvent { sandbox_id, event: Some(event) } + payload: Some(watch_sandboxes_event::Payload::Sandbox( + WatchSandboxesSandboxEvent { sandbox: Some(sandbox) } )), }; if tx.send(Ok(event)).await.is_err() { @@ -1124,14 +1803,71 @@ impl KubernetesComputeDriver { } } } - Ok(Some(Event::Deleted(_))) => {} - Ok(Some(Event::Restarted(_))) => { - debug!(namespace = %namespace, "Kubernetes event watcher restarted"); - } - Ok(None) => { - let _ = tx.send(Err(KubernetesDriverError::Message( - "kubernetes event watcher stream ended".to_string() - ))).await; + Ok(Some(Event::Deleted(obj))) => { + if is_openshell_managed(&obj) + && let Ok(sandbox_id) = sandbox_id_from_object(&obj) + { + remove_indexes(&mut sandbox_name_to_id, &mut agent_pod_to_id, &sandbox_id); + let event = WatchSandboxesEvent { + payload: Some(watch_sandboxes_event::Payload::Deleted( + WatchSandboxesDeletedEvent { sandbox_id } + )), + }; + if tx.send(Ok(event)).await.is_err() { + break; + } + } + } + Ok(Some(Event::Restarted(objs))) => { + for obj in objs { + if let Ok(sandbox) = sandbox_from_claim_object(&namespace, obj) { + let event = WatchSandboxesEvent { + payload: Some(watch_sandboxes_event::Payload::Sandbox( + WatchSandboxesSandboxEvent { sandbox: Some(sandbox) } + )), + }; + if tx.send(Ok(event)).await.is_err() { + return; + } + } + } + } + Ok(None) => { + let _ = tx.send(Err(KubernetesDriverError::Message( + "sandbox claim watcher stream ended unexpectedly".to_string() + ))).await; + break; + } + Err(err) => { + let _ = tx.send(Err(KubernetesDriverError::Message(err.to_string()))).await; + break; + } + }, + result = event_stream.try_next() => match result { + Ok(Some(Event::Applied(obj))) => { + if let Some((sandbox_id, event)) = map_kube_event_to_platform( + &sandbox_name_to_id, + &agent_pod_to_id, + &obj, + ) { + let event = WatchSandboxesEvent { + payload: Some(watch_sandboxes_event::Payload::PlatformEvent( + WatchSandboxesPlatformEvent { sandbox_id, event: Some(event) } + )), + }; + if tx.send(Ok(event)).await.is_err() { + break; + } + } + } + Ok(Some(Event::Deleted(_))) => {} + Ok(Some(Event::Restarted(_))) => { + debug!(namespace = %namespace, "Kubernetes event watcher restarted"); + } + Ok(None) => { + let _ = tx.send(Err(KubernetesDriverError::Message( + "kubernetes event watcher stream ended".to_string() + ))).await; break; } Err(err) => { @@ -1142,9 +1878,1139 @@ impl KubernetesComputeDriver { () = tx.closed() => break, } } - }); + }); + + Ok(Box::pin(ReceiverStream::new(rx))) + } +} + +async fn run_warm_pool_cache_controller( + cache: WarmPoolCache, + client: Client, + watch_client: Client, + fallback_namespace: String, +) { + loop { + refresh_warm_pool_cache(&cache, client.clone(), &fallback_namespace).await; + + let warm_pool_api = match KubernetesComputeDriver::extension_api_for_watch( + watch_client.clone(), + &fallback_namespace, + SANDBOX_WARM_POOL_KIND, + ) + .await + { + Ok(api) => api, + Err(err) => { + cache.mark_not_ready().await; + warn!( + namespace = %fallback_namespace, + error = %err, + "Warm-pool watch is unavailable; warm-pool allocation disabled until retry" + ); + tokio::time::sleep(WARM_POOL_CACHE_RETRY_DELAY).await; + continue; + } + }; + let template_api = match KubernetesComputeDriver::extension_api_for_watch( + watch_client.clone(), + &fallback_namespace, + SANDBOX_TEMPLATE_KIND, + ) + .await + { + Ok(api) => api, + Err(err) => { + cache.mark_not_ready().await; + warn!( + namespace = %fallback_namespace, + error = %err, + "SandboxTemplate watch is unavailable; warm-pool allocation disabled until retry" + ); + tokio::time::sleep(WARM_POOL_CACHE_RETRY_DELAY).await; + continue; + } + }; + + let warm_pool_config = + watcher::Config::default().labels(&format!("{LABEL_WARM_POOL_ENABLED}=true")); + let template_config = watcher::Config::default(); + let mut warm_pool_stream = watcher::watcher(warm_pool_api, warm_pool_config).boxed(); + let mut template_stream = watcher::watcher(template_api, template_config).boxed(); + + loop { + tokio::select! { + event = warm_pool_stream.try_next() => { + if !handle_warm_pool_cache_watch_event( + &cache, + client.clone(), + &fallback_namespace, + event, + SANDBOX_WARM_POOL_KIND, + ) + .await + { + break; + } + } + event = template_stream.try_next() => { + if !handle_warm_pool_cache_watch_event( + &cache, + client.clone(), + &fallback_namespace, + event, + SANDBOX_TEMPLATE_KIND, + ) + .await + { + break; + } + } + } + } + + cache.mark_not_ready().await; + tokio::time::sleep(WARM_POOL_CACHE_RETRY_DELAY).await; + } +} + +async fn handle_warm_pool_cache_watch_event( + cache: &WarmPoolCache, + client: Client, + fallback_namespace: &str, + event: Result>, watcher::Error>, + kind: &str, +) -> bool { + match event { + Ok(Some(Event::Applied(_) | Event::Deleted(_) | Event::Restarted(_))) => { + refresh_warm_pool_cache(cache, client, fallback_namespace).await; + true + } + Ok(None) => { + warn!( + kind, + "Warm-pool cache watch stream ended; warm-pool allocation disabled until retry" + ); + false + } + Err(err) => { + warn!( + kind, + error = %err, + "Warm-pool cache watch failed; warm-pool allocation disabled until retry" + ); + false + } + } +} + +async fn refresh_warm_pool_cache(cache: &WarmPoolCache, client: Client, fallback_namespace: &str) { + let pools = match KubernetesComputeDriver::list_enabled_warm_pools_with_client( + client.clone(), + fallback_namespace, + ) + .await + { + Ok(pools) => pools, + Err(err) => { + cache.mark_not_ready().await; + warn!( + namespace = %fallback_namespace, + error = %err, + "Failed to refresh warm-pool cache; warm-pool allocation disabled until retry" + ); + return; + } + }; + + let mut entries = Vec::new(); + for pool in pools { + let template_fingerprint = + match KubernetesComputeDriver::template_fingerprint_for_warm_pool_with_client( + client.clone(), + &pool, + ) + .await + { + Ok(Some(fingerprint)) => fingerprint, + Ok(None) => { + debug!( + namespace = %pool.namespace, + warm_pool = %pool.name, + template_namespace = %pool.template_namespace, + template = %pool.template_name, + "Skipping warm pool with missing or unreadable SandboxTemplate" + ); + continue; + } + Err(err) => { + warn!( + namespace = %pool.namespace, + warm_pool = %pool.name, + template_namespace = %pool.template_namespace, + template = %pool.template_name, + error = %err, + "Skipping warm pool after SandboxTemplate fingerprint failed" + ); + continue; + } + }; + + debug!( + namespace = %pool.namespace, + warm_pool = %pool.name, + template_namespace = %pool.template_namespace, + template = %pool.template_name, + template_fingerprint = %template_fingerprint, + "Cached OpenShell-enabled warm-pool template fingerprint" + ); + + entries.push(WarmPoolCacheEntry { + namespace: pool.namespace, + name: pool.name, + template_namespace: pool.template_namespace, + template_name: pool.template_name, + template_fingerprint, + }); + } + + let count = entries.len(); + cache.replace_entries(entries).await; + debug!( + namespace = %fallback_namespace, + count, + "Warm-pool cache refreshed" + ); +} + +async fn run_warm_pool_profile_reconciler( + client: Client, + watch_client: Client, + config: KubernetesComputeConfig, +) { + let profile_namespace = config + .warm_pooling + .profiles + .effective_namespace(&config.namespace) + .to_string(); + let label_selector = config + .warm_pooling + .profiles + .effective_label_selector() + .to_string(); + + loop { + reconcile_existing_warm_pool_profiles( + client.clone(), + &config, + &profile_namespace, + &label_selector, + ) + .await; + + let api: Api = Api::namespaced(watch_client.clone(), &profile_namespace); + let watcher_config = watcher::Config::default().labels(&label_selector); + let mut stream = watcher::watcher(api, watcher_config).boxed(); + + loop { + tokio::select! { + event = stream.try_next() => { + match event { + Ok(Some(Event::Applied(config_map))) => { + reconcile_warm_pool_profile_config_map(client.clone(), &config, config_map).await; + } + Ok(Some(Event::Deleted(config_map))) => { + garbage_collect_warm_pool_profile(client.clone(), &config, &config_map).await; + } + Ok(Some(Event::Restarted(config_maps))) => { + for config_map in config_maps { + reconcile_warm_pool_profile_config_map(client.clone(), &config, config_map).await; + } + } + Ok(None) => { + warn!( + namespace = %profile_namespace, + "Warm-pool profile watch stream ended; retrying" + ); + break; + } + Err(err) => { + warn!( + namespace = %profile_namespace, + error = %err, + "Warm-pool profile watch failed; retrying" + ); + break; + } + } + } + () = tokio::time::sleep(WARM_POOL_PROFILE_RESYNC_DELAY) => { + reconcile_existing_warm_pool_profiles( + client.clone(), + &config, + &profile_namespace, + &label_selector, + ) + .await; + } + } + } + + tokio::time::sleep(WARM_POOL_PROFILE_REQUEUE_DELAY).await; + } +} + +async fn resolve_sandbox_identity_for_config( + client: Client, + config: &KubernetesComputeConfig, + namespace: &str, +) -> (u32, u32, BTreeMap) { + // Explicit config takes priority — skip namespace lookup entirely. + if config.sandbox_uid.is_some() { + let uid = config.resolve_sandbox_uid(None); + let gid = config.resolve_sandbox_gid(uid, None); + return (uid, gid, BTreeMap::new()); + } + + let ns_api: Api = Api::all(client); + match tokio::time::timeout(KUBE_API_TIMEOUT, ns_api.get(namespace)).await { + Ok(Ok(ns)) => { + let anns = ns.metadata.annotations.unwrap_or_default(); + tracing::info!( + namespace, + uid_range = ?anns.get(crate::config::ANNOTATION_SCC_UID_RANGE), + sup_groups = ?anns.get(crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS), + "Resolved namespace annotations for sandbox identity" + ); + let uid = config.resolve_sandbox_uid(Some(&anns)); + let baseline_gid = config.resolve_sandbox_gid(uid, None); + let gid = config.sandbox_gid.map_or_else( + || { + anns.get(crate::config::ANNOTATION_SCC_SUPPLEMENTAL_GROUPS) + .and_then(|sup_range| { + KubernetesComputeConfig::from_open_shift_supplemental_groups(sup_range) + }) + .unwrap_or(baseline_gid) + }, + |_| baseline_gid, + ); + tracing::info!(uid, gid, "Resolved sandbox identity"); + (uid, gid, anns) + } + Ok(Err(e)) => { + tracing::warn!( + namespace, + error = %e, + "Failed to fetch namespace for SCC annotations, falling back to defaults" + ); + let uid = DEFAULT_SANDBOX_UID; + let gid = config.resolve_sandbox_gid(uid, None); + (uid, gid, BTreeMap::new()) + } + Err(_) => { + tracing::warn!( + namespace, + "Namespace fetch timed out, falling back to defaults" + ); + let uid = DEFAULT_SANDBOX_UID; + let gid = config.resolve_sandbox_gid(uid, None); + (uid, gid, BTreeMap::new()) + } + } +} + +async fn reconcile_existing_warm_pool_profiles( + client: Client, + config: &KubernetesComputeConfig, + profile_namespace: &str, + label_selector: &str, +) { + let api: Api = Api::namespaced(client.clone(), profile_namespace); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + api.list(&ListParams::default().labels(label_selector)), + ) + .await + { + Ok(Ok(list)) => { + let active_source_uids = list + .items + .iter() + .filter_map(|config_map| config_map.metadata.uid.clone()) + .collect::>(); + for config_map in list.items { + reconcile_warm_pool_profile_config_map(client.clone(), config, config_map).await; + } + if let Err(err) = garbage_collect_orphaned_warm_pool_profile_resources( + client, + &config.namespace, + &active_source_uids, + ) + .await + { + warn!( + namespace = %config.namespace, + error = %err, + "Failed to garbage-collect orphaned warm-pool profile resources" + ); + } + } + Ok(Err(err)) => { + warn!( + namespace = %profile_namespace, + error = %err, + "Failed to list warm-pool profile ConfigMaps" + ); + } + Err(_) => { + warn!( + namespace = %profile_namespace, + "Timed out listing warm-pool profile ConfigMaps" + ); + } + } +} + +async fn reconcile_warm_pool_profile_config_map( + client: Client, + config: &KubernetesComputeConfig, + config_map: ConfigMap, +) { + match render_warm_pool_profile(client.clone(), config, &config_map).await { + Ok(rendered) => { + if let Err(err) = apply_rendered_warm_pool_profile(client.clone(), &rendered).await { + warn!( + namespace = %rendered.source.namespace, + config_map = %rendered.source.name, + error = %err, + "Failed to apply warm-pool profile" + ); + patch_warm_pool_profile_status( + client, + &rendered.source.namespace, + &rendered.source.name, + "Error", + Some(&rendered), + &err, + ) + .await; + return; + } + if let Err(err) = + garbage_collect_superseded_warm_pool_profile_resources(client.clone(), &rendered) + .await + { + warn!( + namespace = %rendered.source.namespace, + config_map = %rendered.source.name, + error = %err, + "Failed to garbage-collect superseded warm-pool profile resources" + ); + } + patch_warm_pool_profile_status( + client, + &rendered.source.namespace, + &rendered.source.name, + "Ready", + Some(&rendered), + "", + ) + .await; + } + Err(err) => { + let namespace = config_map + .metadata + .namespace + .as_deref() + .unwrap_or_else(|| { + config + .warm_pooling + .profiles + .effective_namespace(&config.namespace) + }) + .to_string(); + let Some(name) = config_map.metadata.name.as_deref() else { + warn!(error = %err, "Warm-pool profile ConfigMap is missing a name"); + return; + }; + warn!( + namespace = %namespace, + config_map = %name, + error = %err, + "Warm-pool profile is invalid" + ); + patch_warm_pool_profile_status(client, &namespace, name, "Error", None, &err).await; + } + } +} + +async fn render_warm_pool_profile( + client: Client, + config: &KubernetesComputeConfig, + config_map: &ConfigMap, +) -> Result { + let source = warm_pool_profile_source(config, config_map)?; + let data = config_map + .data + .as_ref() + .and_then(|data| data.get(WARM_POOL_PROFILE_DATA_KEY)) + .ok_or_else(|| format!("missing data key '{WARM_POOL_PROFILE_DATA_KEY}'"))?; + let profile = parse_warm_pool_profile(data)?; + validate_warm_pool_profile( + &profile, + config.warm_pooling.profiles.effective_max_replicas(), + )?; + + let target_namespace = config.namespace.clone(); + let spec = warm_pool_profile_to_sandbox_spec(&profile); + let (profile_user_id, profile_group_id, _) = + resolve_sandbox_identity_for_config(client, config, &target_namespace).await; + let params = SandboxPodParams { + default_image: &config.default_image, + image_pull_policy: &config.image_pull_policy, + image_pull_secrets: &config.image_pull_secrets, + supervisor_image: &config.supervisor_image, + supervisor_image_pull_policy: &config.supervisor_image_pull_policy, + supervisor_sideload_method: config.supervisor_sideload_method, + topology: config.topology, + proxy_uid: config.sidecar.proxy_uid, + process_binary_aware_network_policy: config.sidecar.process_binary_aware_network_policy, + service_account_name: &config.service_account_name, + sandbox_id: "", + sandbox_name: "", + grpc_endpoint: &config.grpc_endpoint, + ssh_socket_path: &config.ssh_socket_path, + client_tls_secret_name: &config.client_tls_secret_name, + host_gateway_ip: &config.host_gateway_ip, + enable_user_namespaces: config.enable_user_namespaces, + app_armor_profile: config.app_armor_profile.as_ref(), + workspace_default_storage_size: &config.workspace_default_storage_size, + default_runtime_class_name: &config.default_runtime_class_name, + sa_token_ttl_secs: config.effective_sa_token_ttl_secs(), + provider_spiffe_enabled: config.provider_spiffe_enabled(), + provider_spiffe_workload_api_socket_path: &config.provider_spiffe_workload_api_socket_path, + sandbox_uid: profile_user_id, + sandbox_gid: profile_group_id, + }; + validate_sidecar_proxy_identity(¶ms).map_err(|err| err.to_string())?; + + let mut rendered = sandbox_to_k8s_spec(Some(&spec), ¶ms)?; + remove_warm_pool_template_identity_env(&mut rendered); + ensure_pod_dns_policy(&mut rendered); + let fingerprint = sandbox_spec_fingerprint(&rendered)?; + let generated_name = generated_warm_pool_profile_name(&source.name, &fingerprint); + validate_kubernetes_dns1123_label(&generated_name, "generated warm-pool resource name")?; + let template_spec = rendered + .get("spec") + .cloned() + .ok_or_else(|| "rendered sandbox spec is missing spec".to_string())?; + + Ok(RenderedWarmPoolProfile { + source, + workspace: profile.workspace, + target_namespace, + generated_name, + replicas: profile.replicas, + template_spec, + fingerprint, + }) +} + +fn warm_pool_profile_source( + config: &KubernetesComputeConfig, + config_map: &ConfigMap, +) -> Result { + let namespace = config_map.metadata.namespace.clone().unwrap_or_else(|| { + config + .warm_pooling + .profiles + .effective_namespace(&config.namespace) + .to_string() + }); + let name = config_map + .metadata + .name + .clone() + .ok_or_else(|| "ConfigMap metadata.name is required".to_string())?; + let uid = config_map + .metadata + .uid + .clone() + .ok_or_else(|| "ConfigMap metadata.uid is required".to_string())?; + let resource_version = config_map + .metadata + .resource_version + .clone() + .unwrap_or_default(); + Ok(WarmPoolProfileSource { + namespace, + name, + uid, + resource_version, + }) +} + +fn parse_warm_pool_profile(data: &str) -> Result { + toml::from_str::(data) + .map_err(|err| format!("failed to parse {WARM_POOL_PROFILE_DATA_KEY}: {err}")) +} + +fn validate_warm_pool_profile(profile: &WarmPoolProfile, max_replicas: u32) -> Result<(), String> { + if profile.version != WARM_POOL_PROFILE_VERSION { + return Err(format!( + "unsupported warm-pool profile version {}; expected {WARM_POOL_PROFILE_VERSION}", + profile.version + )); + } + if profile.workspace.trim().is_empty() { + return Err("workspace must not be empty".to_string()); + } + if profile.workspace.chars().any(char::is_control) { + return Err("workspace must not contain control characters".to_string()); + } + if profile.replicas == 0 { + return Err("replicas must be greater than 0".to_string()); + } + if profile.replicas > max_replicas { + return Err(format!( + "replicas exceeds maximum ({} > {max_replicas})", + profile.replicas + )); + } + validate_warm_pool_profile_environment(&profile.environment)?; + if let Some(count) = profile.resources.gpu_count + && count == 0 + { + return Err("resources.gpu_count must be greater than 0".to_string()); + } + Ok(()) +} + +fn validate_warm_pool_profile_environment(env: &BTreeMap) -> Result<(), String> { + for (key, value) in env { + if !is_valid_env_key(key) { + return Err(format!( + "environment keys must match ^[A-Za-z_][A-Za-z0-9_]*$; got '{key}'" + )); + } + if key.starts_with("OPENSHELL_") { + return Err(format!( + "environment keys starting with OPENSHELL_ are reserved; got '{key}'" + )); + } + if value.chars().any(char::is_control) { + return Err(format!( + "environment value for '{key}' must not contain control characters" + )); + } + } + Ok(()) +} + +fn is_valid_env_key(key: &str) -> bool { + let mut chars = key.chars(); + let Some(first) = chars.next() else { + return false; + }; + (first == '_' || first.is_ascii_alphabetic()) + && chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric()) +} + +fn warm_pool_profile_to_sandbox_spec(profile: &WarmPoolProfile) -> SandboxSpec { + let resources = (!profile.resources.cpu.is_empty() || !profile.resources.memory.is_empty()) + .then(|| DriverResourceRequirements { + cpu_limit: profile.resources.cpu.clone(), + memory_limit: profile.resources.memory.clone(), + ..Default::default() + }); + + SandboxSpec { + environment: profile.environment.clone().into_iter().collect(), + template: Some(SandboxTemplate { + image: profile.image.clone(), + platform_config: platform_config_for_warm_pool_profile(profile), + resources, + ..Default::default() + }), + resource_requirements: profile + .resources + .gpu_count + .map(|count| ResourceRequirements { + gpu: Some(GpuResourceRequirements { count: Some(count) }), + }), + ..Default::default() + } +} + +fn platform_config_for_warm_pool_profile(profile: &WarmPoolProfile) -> Option { + use prost_types::{Struct, Value, value::Kind}; + + if profile.runtime_class_name.is_empty() { + return None; + } + + Some(Struct { + fields: BTreeMap::from([( + "runtime_class_name".to_string(), + Value { + kind: Some(Kind::StringValue(profile.runtime_class_name.clone())), + }, + )]), + }) +} + +fn ensure_pod_dns_policy(rendered: &mut serde_json::Value) { + if let Some(spec) = rendered + .pointer_mut("/spec/podTemplate/spec") + .and_then(serde_json::Value::as_object_mut) + { + spec.entry("dnsPolicy".to_string()) + .or_insert_with(|| serde_json::json!("ClusterFirst")); + } +} + +fn remove_warm_pool_template_identity_env(rendered: &mut serde_json::Value) { + if let Some(containers) = rendered + .pointer_mut("/spec/podTemplate/spec/containers") + .and_then(serde_json::Value::as_array_mut) + { + for container in containers { + remove_sandbox_identity_env(container); + } + } +} + +fn generated_warm_pool_profile_name(profile_name: &str, fingerprint: &str) -> String { + let hash = fingerprint + .strip_prefix("sha256:") + .unwrap_or(fingerprint) + .chars() + .filter(char::is_ascii_hexdigit) + .take(8) + .collect::() + .to_ascii_lowercase(); + let suffix = if hash.is_empty() { + "00000000".to_string() + } else { + hash + }; + let prefix = sanitize_dns_label_segment(profile_name); + let reserved = WARM_POOL_PROFILE_NAME_PREFIX.len() + 1 + 1 + suffix.len(); + let max_prefix_len = MAX_KUBE_NAME_LEN.saturating_sub(reserved); + let mut trimmed = prefix.chars().take(max_prefix_len).collect::(); + trimmed = trimmed.trim_matches('-').to_string(); + if trimmed.is_empty() { + trimmed = "profile".to_string(); + } + format!("{WARM_POOL_PROFILE_NAME_PREFIX}-{trimmed}-{suffix}") +} + +fn sanitize_dns_label_segment(value: &str) -> String { + let mut out = String::new(); + let mut last_dash = false; + for byte in value.bytes() { + let ch = if byte.is_ascii_lowercase() || byte.is_ascii_digit() { + byte as char + } else if byte.is_ascii_uppercase() { + (byte as char).to_ascii_lowercase() + } else { + '-' + }; + if ch == '-' { + if !last_dash { + out.push(ch); + } + last_dash = true; + } else { + out.push(ch); + last_dash = false; + } + } + let trimmed = out.trim_matches('-').to_string(); + if trimmed.is_empty() { + "profile".to_string() + } else { + trimmed + } +} + +async fn apply_rendered_warm_pool_profile( + client: Client, + rendered: &RenderedWarmPoolProfile, +) -> Result<(), String> { + let template_api = + KubernetesComputeDriver::sandbox_template_api(client.clone(), &rendered.target_namespace); + let warm_pool_api = + KubernetesComputeDriver::sandbox_warm_pool_api(client.clone(), &rendered.target_namespace); + let template = rendered_warm_pool_template_object(rendered, &template_api.resource); + let warm_pool = rendered_warm_pool_object(rendered, &warm_pool_api.resource); + apply_dynamic_object(&template_api.api, &rendered.generated_name, &template).await?; + apply_dynamic_object(&warm_pool_api.api, &rendered.generated_name, &warm_pool).await?; + Ok(()) +} + +async fn apply_dynamic_object( + api: &Api, + name: &str, + obj: &DynamicObject, +) -> Result<(), String> { + match tokio::time::timeout(KUBE_API_TIMEOUT, api.create(&PostParams::default(), obj)).await { + Ok(Ok(_)) => Ok(()), + Ok(Err(KubeError::Api(err))) if err.code == 409 => { + let patch = dynamic_object_merge_patch(obj); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + api.patch(name, &PatchParams::default(), &Patch::Merge(&patch)), + ) + .await + { + Ok(Ok(_)) => Ok(()), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } + } + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } +} - Ok(Box::pin(ReceiverStream::new(rx))) +fn dynamic_object_merge_patch(obj: &DynamicObject) -> serde_json::Value { + serde_json::json!({ + "metadata": { + "labels": obj.metadata.labels, + "annotations": obj.metadata.annotations, + }, + "spec": obj.data.get("spec").cloned().unwrap_or_else(|| serde_json::json!({})), + }) +} + +fn rendered_warm_pool_template_object( + rendered: &RenderedWarmPoolProfile, + resource: &ApiResource, +) -> DynamicObject { + let mut obj = DynamicObject::new(&rendered.generated_name, resource); + obj.metadata = ObjectMeta { + name: Some(rendered.generated_name.clone()), + namespace: Some(rendered.target_namespace.clone()), + labels: Some(warm_pool_profile_generated_labels(rendered)), + annotations: Some(warm_pool_profile_generated_annotations(rendered)), + ..Default::default() + }; + obj.data = serde_json::json!({ + "spec": rendered.template_spec, + }); + obj +} + +fn rendered_warm_pool_object( + rendered: &RenderedWarmPoolProfile, + resource: &ApiResource, +) -> DynamicObject { + let mut obj = DynamicObject::new(&rendered.generated_name, resource); + obj.metadata = ObjectMeta { + name: Some(rendered.generated_name.clone()), + namespace: Some(rendered.target_namespace.clone()), + labels: Some(warm_pool_profile_generated_labels(rendered)), + annotations: Some(warm_pool_profile_generated_annotations(rendered)), + ..Default::default() + }; + obj.data = serde_json::json!({ + "spec": { + "replicas": rendered.replicas, + "sandboxTemplateRef": { + "name": rendered.generated_name + } + } + }); + obj +} + +fn warm_pool_profile_generated_labels( + rendered: &RenderedWarmPoolProfile, +) -> BTreeMap { + BTreeMap::from([ + (LABEL_WARM_POOL_ENABLED.to_string(), "true".to_string()), + ( + LABEL_WARM_POOL_MANAGED_BY.to_string(), + LABEL_WARM_POOL_MANAGED_BY_VALUE.to_string(), + ), + ( + LABEL_WARM_POOL_PROFILE.to_string(), + label_value_for_profile_name(&rendered.source.name), + ), + ( + LABEL_WARM_POOL_PROFILE_UID.to_string(), + rendered.source.uid.clone(), + ), + ]) +} + +fn label_value_for_profile_name(name: &str) -> String { + let sanitized = sanitize_dns_label_segment(name); + let mut value = sanitized.chars().take(63).collect::(); + value = value.trim_matches('-').to_string(); + if value.is_empty() { + "profile".to_string() + } else { + value + } +} + +fn warm_pool_profile_generated_annotations( + rendered: &RenderedWarmPoolProfile, +) -> BTreeMap { + BTreeMap::from([ + ( + ANNOTATION_WARM_POOL_PROFILE_NAMESPACE.to_string(), + rendered.source.namespace.clone(), + ), + ( + ANNOTATION_WARM_POOL_PROFILE_NAME.to_string(), + rendered.source.name.clone(), + ), + ( + ANNOTATION_WARM_POOL_PROFILE_UID.to_string(), + rendered.source.uid.clone(), + ), + ( + ANNOTATION_WARM_POOL_SOURCE_RESOURCE_VERSION.to_string(), + rendered.source.resource_version.clone(), + ), + ( + ANNOTATION_WARM_POOL_TEMPLATE_FINGERPRINT.to_string(), + rendered.fingerprint.clone(), + ), + ( + LABEL_SANDBOX_WORKSPACE.to_string(), + rendered.workspace.clone(), + ), + ]) +} + +async fn garbage_collect_superseded_warm_pool_profile_resources( + client: Client, + rendered: &RenderedWarmPoolProfile, +) -> Result<(), String> { + garbage_collect_warm_pool_profile_resources( + client, + &rendered.target_namespace, + &rendered.source.uid, + Some(&rendered.generated_name), + ) + .await +} + +async fn garbage_collect_warm_pool_profile( + client: Client, + config: &KubernetesComputeConfig, + cm: &ConfigMap, +) { + let Some(uid) = cm.metadata.uid.as_deref() else { + return; + }; + let namespace = config.namespace.clone(); + if let Err(err) = + garbage_collect_warm_pool_profile_resources(client, &namespace, uid, None).await + { + warn!( + namespace = %namespace, + config_map_uid = %uid, + error = %err, + "Failed to garbage-collect deleted warm-pool profile resources" + ); + } +} + +async fn garbage_collect_warm_pool_profile_resources( + client: Client, + target_namespace: &str, + source_uid: &str, + keep_name: Option<&str>, +) -> Result<(), String> { + let selector = format!( + "{LABEL_WARM_POOL_MANAGED_BY}={LABEL_WARM_POOL_MANAGED_BY_VALUE},{LABEL_WARM_POOL_PROFILE_UID}={source_uid}" + ); + let lp = ListParams::default().labels(&selector); + let warm_pool_api = + KubernetesComputeDriver::sandbox_warm_pool_api(client.clone(), target_namespace); + delete_matching_dynamic_objects(&warm_pool_api.api, &lp, keep_name).await?; + let template_api = KubernetesComputeDriver::sandbox_template_api(client, target_namespace); + delete_matching_dynamic_objects(&template_api.api, &lp, keep_name).await?; + Ok(()) +} + +async fn garbage_collect_orphaned_warm_pool_profile_resources( + client: Client, + target_namespace: &str, + active_source_uids: &HashSet, +) -> Result<(), String> { + let selector = format!("{LABEL_WARM_POOL_MANAGED_BY}={LABEL_WARM_POOL_MANAGED_BY_VALUE}"); + let lp = ListParams::default().labels(&selector); + let warm_pool_api = + KubernetesComputeDriver::sandbox_warm_pool_api(client.clone(), target_namespace); + delete_orphaned_dynamic_objects(&warm_pool_api.api, &lp, active_source_uids).await?; + let template_api = KubernetesComputeDriver::sandbox_template_api(client, target_namespace); + delete_orphaned_dynamic_objects(&template_api.api, &lp, active_source_uids).await?; + Ok(()) +} + +async fn delete_orphaned_dynamic_objects( + api: &Api, + lp: &ListParams, + active_source_uids: &HashSet, +) -> Result<(), String> { + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, api.list(lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + for obj in list.items { + if !warm_pool_profile_generated_resource_is_orphaned(&obj, active_source_uids) { + continue; + } + let Some(name) = obj.metadata.name.as_deref() else { + continue; + }; + match tokio::time::timeout(KUBE_API_TIMEOUT, api.delete(name, &DeleteParams::default())) + .await + { + Ok(Ok(_)) => {} + Ok(Err(KubeError::Api(err))) if err.code == 404 => {} + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(()) +} + +fn warm_pool_profile_generated_resource_is_orphaned( + obj: &DynamicObject, + active_source_uids: &HashSet, +) -> bool { + let labels = obj.metadata.labels.as_ref(); + let managed_by_driver = labels + .and_then(|labels| labels.get(LABEL_WARM_POOL_MANAGED_BY)) + .is_some_and(|value| value == LABEL_WARM_POOL_MANAGED_BY_VALUE); + let Some(source_uid) = labels.and_then(|labels| labels.get(LABEL_WARM_POOL_PROFILE_UID)) else { + return false; + }; + managed_by_driver && !active_source_uids.contains(source_uid) +} + +async fn delete_matching_dynamic_objects( + api: &Api, + lp: &ListParams, + keep_name: Option<&str>, +) -> Result<(), String> { + let list = match tokio::time::timeout(KUBE_API_TIMEOUT, api.list(lp)).await { + Ok(Ok(list)) => list, + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + for obj in list.items { + let Some(name) = obj.metadata.name.as_deref() else { + continue; + }; + if keep_name == Some(name) { + continue; + } + match tokio::time::timeout(KUBE_API_TIMEOUT, api.delete(name, &DeleteParams::default())) + .await + { + Ok(Ok(_)) => {} + Ok(Err(KubeError::Api(err))) if err.code == 404 => {} + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(()) +} + +async fn patch_warm_pool_profile_status( + client: Client, + namespace: &str, + name: &str, + status: &str, + rendered: Option<&RenderedWarmPoolProfile>, + error: &str, +) { + let api: Api = Api::namespaced(client, namespace); + let mut annotations = serde_json::Map::new(); + annotations.insert( + ANNOTATION_WARM_POOL_STATUS.to_string(), + serde_json::json!(status), + ); + annotations.insert( + ANNOTATION_WARM_POOL_LAST_ERROR.to_string(), + serde_json::json!(error), + ); + if let Some(rendered) = rendered { + annotations.insert( + ANNOTATION_WARM_POOL_OBSERVED_RESOURCE_VERSION.to_string(), + serde_json::json!(rendered.source.resource_version), + ); + annotations.insert( + ANNOTATION_WARM_POOL_TARGET_NAMESPACE.to_string(), + serde_json::json!(rendered.target_namespace), + ); + annotations.insert( + ANNOTATION_WARM_POOL_TEMPLATE.to_string(), + serde_json::json!(rendered.generated_name), + ); + annotations.insert( + ANNOTATION_WARM_POOL_FINGERPRINT.to_string(), + serde_json::json!(rendered.fingerprint), + ); + } + let patch = serde_json::json!({ + "metadata": { + "annotations": annotations + } + }); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + api.patch(name, &PatchParams::default(), &Patch::Merge(&patch)), + ) + .await + { + Ok(Ok(_)) => {} + Ok(Err(err)) => warn!( + namespace, + config_map = name, + error = %err, + "Failed to patch warm-pool profile status annotations" + ), + Err(_) => warn!( + namespace, + config_map = name, + "Timed out patching warm-pool profile status annotations" + ), } } @@ -1196,15 +3062,356 @@ fn sandbox_labels(sandbox: &Sandbox) -> BTreeMap { labels } -fn sandbox_annotations(sandbox: &Sandbox) -> BTreeMap { - let mut annotations = BTreeMap::new(); - annotations.insert(LABEL_SANDBOX_ID.to_string(), sandbox.id.clone()); - annotations.insert(LABEL_SANDBOX_NAME.to_string(), sandbox.name.clone()); - annotations.insert( - LABEL_SANDBOX_WORKSPACE.to_string(), - sandbox.workspace.clone(), - ); - annotations +fn sandbox_annotations(sandbox: &Sandbox) -> BTreeMap { + let mut annotations = BTreeMap::new(); + annotations.insert(LABEL_SANDBOX_ID.to_string(), sandbox.id.clone()); + annotations.insert(LABEL_SANDBOX_NAME.to_string(), sandbox.name.clone()); + annotations.insert( + LABEL_SANDBOX_WORKSPACE.to_string(), + sandbox.workspace.clone(), + ); + annotations +} + +fn sandbox_claim_to_k8s_object( + sandbox: &Sandbox, + warm_pool_name: &str, + resource: &ApiResource, +) -> DynamicObject { + let kube_name = kube_resource_name(&sandbox.workspace, &sandbox.name); + let mut obj = DynamicObject::new(&kube_name, resource); + let mut labels = sandbox_labels(sandbox); + labels.insert( + LABEL_ALLOCATION.to_string(), + LABEL_ALLOCATION_SANDBOX_CLAIM.to_string(), + ); + obj.metadata = ObjectMeta { + name: Some(kube_name), + labels: Some(labels), + annotations: Some(sandbox_annotations(sandbox)), + ..Default::default() + }; + obj.data = serde_json::json!({ + "spec": { + "warmPoolRef": { + "name": warm_pool_name + } + } + }); + obj +} + +fn enabled_warm_pool_from_object(obj: DynamicObject) -> Option { + let namespace = obj.metadata.namespace.clone()?; + let name = obj.metadata.name.clone()?; + let template_name = string_at(&obj.data, &["spec", "sandboxTemplateRef", "name"])?; + let template_namespace = string_at(&obj.data, &["spec", "sandboxTemplateRef", "namespace"]) + .unwrap_or_else(|| namespace.clone()); + Some(EnabledWarmPool { + namespace, + name, + template_namespace, + template_name, + }) +} + +fn sandbox_template_fingerprint(obj: &DynamicObject) -> Result { + sandbox_spec_fingerprint(&obj.data) +} + +fn sandbox_spec_fingerprint(data: &serde_json::Value) -> Result { + let spec = data + .get("spec") + .ok_or_else(|| "object is missing spec".to_string())?; + let mut normalized = spec.clone(); + normalize_template_spec_for_fingerprint(&mut normalized); + stable_json_fingerprint(&normalized) +} + +fn sandbox_spec_fingerprint_debug_json(data: &serde_json::Value) -> Result { + let spec = data + .get("spec") + .ok_or_else(|| "object is missing spec".to_string())?; + let mut normalized = spec.clone(); + normalize_template_spec_for_fingerprint(&mut normalized); + serde_json::to_string(&canonical_json_value(&normalized)) + .map_err(|err| format!("failed to serialize normalized template: {err}")) +} + +fn normalize_template_spec_for_fingerprint(value: &mut serde_json::Value) { + remove_path_if_value( + value, + &["envVarsInjectionPolicy"], + &serde_json::json!("Disallowed"), + ); + remove_path_if_value( + value, + &["networkPolicyManagement"], + &serde_json::json!("Managed"), + ); + remove_path_if_value(value, &["operatingMode"], &serde_json::json!("Running")); + remove_path_if_value(value, &["replicas"], &serde_json::json!(1)); + remove_path_if_value(value, &["shutdownPolicy"], &serde_json::json!("Retain")); + remove_path_if_value( + value, + &["podTemplate", "spec", "dnsPolicy"], + &serde_json::json!("ClusterFirst"), + ); + remove_path_if_value( + value, + &["volumeClaimTemplatesPolicy"], + &serde_json::json!("Disallowed"), + ); + remove_path( + value, + &[ + "podTemplate", + "metadata", + "annotations", + POD_ANNOTATION_SANDBOX_ID, + ], + ); + remove_path( + value, + &["podTemplate", "metadata", "labels", LABEL_SANDBOX_ID], + ); + remove_empty_object_path(value, &["podTemplate", "metadata", "annotations"]); + remove_empty_object_path(value, &["podTemplate", "metadata", "labels"]); + remove_empty_object_path(value, &["podTemplate", "metadata"]); + if let Some(containers) = value + .pointer_mut("/podTemplate/spec/containers") + .and_then(serde_json::Value::as_array_mut) + { + for container in containers { + remove_sandbox_identity_env(container); + remove_empty_container_resources(container); + remove_default_volume_mount_read_only(container); + } + } + if let Some(init_containers) = value + .pointer_mut("/podTemplate/spec/initContainers") + .and_then(serde_json::Value::as_array_mut) + { + for container in init_containers { + remove_empty_container_resources(container); + remove_default_volume_mount_read_only(container); + } + } +} + +fn remove_sandbox_identity_env(container: &mut serde_json::Value) { + let Some(env) = container + .get_mut("env") + .and_then(serde_json::Value::as_array_mut) + else { + return; + }; + env.retain(|entry| { + !matches!( + entry.get("name").and_then(serde_json::Value::as_str), + Some(name) + if name == openshell_core::sandbox_env::SANDBOX_ID + || name == openshell_core::sandbox_env::SANDBOX + ) + }); +} + +fn remove_path(value: &mut serde_json::Value, path: &[&str]) { + let Some((last, parents)) = path.split_last() else { + return; + }; + let mut current = value; + for key in parents { + let Some(next) = current.get_mut(*key) else { + return; + }; + current = next; + } + if let Some(object) = current.as_object_mut() { + object.remove(*last); + } +} + +fn remove_path_if_value( + value: &mut serde_json::Value, + path: &[&str], + expected: &serde_json::Value, +) { + let Some((last, parents)) = path.split_last() else { + return; + }; + let mut current = value; + for key in parents { + let Some(next) = current.get_mut(*key) else { + return; + }; + current = next; + } + if current.get(*last) == Some(expected) + && let Some(object) = current.as_object_mut() + { + object.remove(*last); + } +} + +fn remove_empty_container_resources(container: &mut serde_json::Value) { + if container + .get("resources") + .is_some_and(|resources| resources.as_object().is_some_and(serde_json::Map::is_empty)) + && let Some(object) = container.as_object_mut() + { + object.remove("resources"); + } +} + +fn remove_default_volume_mount_read_only(container: &mut serde_json::Value) { + let Some(volume_mounts) = container + .get_mut("volumeMounts") + .and_then(serde_json::Value::as_array_mut) + else { + return; + }; + for mount in volume_mounts { + remove_path_if_value(mount, &["readOnly"], &serde_json::json!(false)); + } +} + +fn remove_empty_object_path(value: &mut serde_json::Value, path: &[&str]) { + let Some((last, parents)) = path.split_last() else { + return; + }; + let mut current = value; + for key in parents { + let Some(next) = current.get_mut(*key) else { + return; + }; + current = next; + } + if current + .get(*last) + .is_some_and(|entry| entry.as_object().is_some_and(serde_json::Map::is_empty)) + && let Some(object) = current.as_object_mut() + { + object.remove(*last); + } +} + +fn stable_json_fingerprint(value: &serde_json::Value) -> Result { + let canonical = canonical_json_value(value); + let bytes = serde_json::to_vec(&canonical) + .map_err(|err| format!("failed to serialize canonical template: {err}"))?; + let digest = Sha256::digest(bytes); + Ok(hex_encode(&digest)) +} + +fn canonical_json_value(value: &serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Array(items) => { + serde_json::Value::Array(items.iter().map(canonical_json_value).collect()) + } + serde_json::Value::Object(object) => { + let mut sorted = serde_json::Map::new(); + let mut keys = object.keys().collect::>(); + keys.sort(); + for key in keys { + if let Some(value) = object.get(key) { + sorted.insert(key.clone(), canonical_json_value(value)); + } + } + serde_json::Value::Object(sorted) + } + other => other.clone(), + } +} + +fn hex_encode(bytes: &[u8]) -> String { + const HEX: &[u8; 16] = b"0123456789abcdef"; + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes { + out.push(HEX[(byte >> 4) as usize] as char); + out.push(HEX[(byte & 0x0f) as usize] as char); + } + out +} + +fn string_at(data: &serde_json::Value, path: &[&str]) -> Option { + let mut current = data; + for key in path { + current = current.get(*key)?; + } + current + .as_str() + .filter(|value| !value.is_empty()) + .map(ToString::to_string) +} + +fn sandbox_from_claim_object( + default_namespace: &str, + obj: DynamicObject, +) -> Result { + let kube_name = obj.metadata.name.clone().unwrap_or_default(); + if !is_openshell_managed(&obj) { + debug!(object = %kube_name, "skipping SandboxClaim not managed by openshell"); + return Err(format!("SandboxClaim {kube_name} not managed by openshell")); + } + let id = sandbox_id_from_object(&obj)?; + let Some(name) = annotation_or_label(&obj, LABEL_SANDBOX_NAME) else { + return Err(format!("SandboxClaim {kube_name} missing sandbox name")); + }; + let Some(workspace) = annotation_or_label(&obj, LABEL_SANDBOX_WORKSPACE) else { + return Err(format!( + "SandboxClaim {kube_name} missing sandbox workspace" + )); + }; + let namespace = obj + .metadata + .namespace + .clone() + .unwrap_or_else(|| default_namespace.to_string()); + Ok(Sandbox { + id, + name, + namespace, + spec: None, + status: Some(claim_status_from_object(&obj)), + workspace, + }) +} + +fn claim_status_from_object(obj: &DynamicObject) -> SandboxStatus { + let status_obj = obj + .data + .get("status") + .and_then(serde_json::Value::as_object); + let sandbox_name = status_obj + .and_then(|status| status.get("sandbox")) + .and_then(|value| value.get("name")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string(); + let conditions = status_obj + .and_then(|status| status.get("conditions")) + .and_then(serde_json::Value::as_array) + .map(|items| { + items + .iter() + .filter_map(condition_from_value) + .collect::>() + }) + .unwrap_or_default(); + + SandboxStatus { + sandbox_name, + instance_id: String::new(), + agent_fd: String::new(), + sandbox_fd: String::new(), + conditions, + deleting: obj.metadata.deletion_timestamp.is_some(), + } +} + +fn extension_api_unavailable(err: &KubeError) -> bool { + matches!(err, KubeError::Api(api) if api.code == 404) } fn sandbox_id_from_object(obj: &DynamicObject) -> Result { @@ -2199,7 +4406,7 @@ struct SandboxPodParams<'a> { workspace_default_storage_size: &'a str, default_runtime_class_name: &'a str, /// Lifetime (seconds) of the projected `ServiceAccount` token used - /// for the bootstrap `IssueSandboxToken` exchange. + /// for the bootstrap `RegisterSupervisorPod` stream. sa_token_ttl_secs: i64, provider_spiffe_enabled: bool, provider_spiffe_workload_api_socket_path: &'a str, @@ -2425,7 +4632,7 @@ fn sandbox_template_to_k8s_with_validated_config( } // Carry the sandbox UUID as a pod annotation so the gateway can resolve // a projected SA token claim (pod name + uid) back to a sandbox identity - // when the supervisor calls `IssueSandboxToken` at startup. The gateway + // when the supervisor calls `RegisterSupervisorPod` at startup. The gateway // also verifies the pod's controlling Sandbox ownerReference against the // live CR before accepting this annotation. Its K8s Role does NOT grant // `patch pods`, so this annotation is effectively immutable post-create. @@ -2633,9 +4840,10 @@ fn sandbox_template_to_k8s_with_validated_config( } // Projected ServiceAccountToken volume — kubelet writes a short-lived // audience-bound JWT into /var/run/secrets/openshell/token and rotates - // it automatically. The supervisor exchanges this for a gateway-minted - // JWT via `IssueSandboxToken` once at startup. In sidecar topology both - // supervisor containers run with the sandbox GID and need group-read access. + // it automatically. The supervisor presents this to `RegisterSupervisorPod` + // and receives a gateway-minted JWT on the activation stream. In sidecar + // topology both supervisor containers run with the sandbox GID and need + // group-read access. let sa_token_default_mode = match params.topology { SupervisorTopology::Combined => 0o400, SupervisorTopology::Sidecar => 0o440, @@ -2975,7 +5183,7 @@ fn apply_required_env( } // Projected ServiceAccount token written by kubelet (see the volume // definition in `sandbox_template_to_k8s`). The supervisor reads this - // and exchanges it for a gateway-minted JWT via `IssueSandboxToken`. + // and receives a gateway-minted JWT via `RegisterSupervisorPod`. upsert_env( env, openshell_core::sandbox_env::K8S_SA_TOKEN_FILE, @@ -3179,6 +5387,492 @@ mod tests { assert!(!should_try_next_sandbox_api_version(&err)); } + #[test] + fn sandbox_claim_object_renders_v1beta1_warm_pool_ref() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_CLAIM_KIND); + let sandbox = Sandbox { + id: "sandbox-id".to_string(), + name: "dev".to_string(), + namespace: String::new(), + workspace: "workspace-a".to_string(), + spec: None, + status: None, + }; + + let claim = sandbox_claim_to_k8s_object(&sandbox, "pool-a", &resource); + + assert_eq!(claim.metadata.name.as_deref(), Some("workspace-a--dev")); + assert_eq!( + claim + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(LABEL_ALLOCATION)) + .map(String::as_str), + Some(LABEL_ALLOCATION_SANDBOX_CLAIM) + ); + assert_eq!( + string_at(&claim.data, &["spec", "warmPoolRef", "name"]).as_deref(), + Some("pool-a") + ); + assert_eq!( + string_at(&claim.data, &["spec", "sandboxTemplateRef", "name"]), + None + ); + } + + #[test] + fn enabled_warm_pool_parse_defaults_template_namespace_to_pool_namespace() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_WARM_POOL_KIND); + let mut warm_pool = DynamicObject::new("pool-a", &resource); + warm_pool.metadata.namespace = Some("team-a".to_string()); + warm_pool.metadata.labels = Some(BTreeMap::from([( + LABEL_WARM_POOL_ENABLED.to_string(), + "true".to_string(), + )])); + warm_pool.data = serde_json::json!({ + "spec": { + "sandboxTemplateRef": { + "name": "template-a" + } + } + }); + + let parsed = enabled_warm_pool_from_object(warm_pool).unwrap(); + + assert_eq!(parsed.namespace, "team-a"); + assert_eq!(parsed.name, "pool-a"); + assert_eq!(parsed.template_namespace, "team-a"); + assert_eq!(parsed.template_name, "template-a"); + } + + fn warm_pool_cache_entry( + namespace: &str, + name: &str, + template_fingerprint: &str, + ) -> WarmPoolCacheEntry { + WarmPoolCacheEntry { + namespace: namespace.to_string(), + name: name.to_string(), + template_namespace: namespace.to_string(), + template_name: format!("{name}-template"), + template_fingerprint: template_fingerprint.to_string(), + } + } + + #[test] + fn warm_pool_cache_requires_successful_refresh_before_matching() { + futures::executor::block_on(async { + let cache = WarmPoolCache::default(); + + assert_eq!( + cache.matching_pool("team-a", "fingerprint-a").await, + WarmPoolCacheLookup::NotReady + ); + + cache.replace_entries(Vec::new()).await; + + assert_eq!( + cache.matching_pool("team-a", "fingerprint-a").await, + WarmPoolCacheLookup::NoMatch + ); + }); + } + + #[test] + fn warm_pool_cache_matches_by_namespace_and_fingerprint() { + futures::executor::block_on(async { + let cache = WarmPoolCache::default(); + cache + .replace_entries(vec![ + warm_pool_cache_entry("team-a", "pool-a", "fingerprint-a"), + warm_pool_cache_entry("team-b", "pool-b", "fingerprint-a"), + ]) + .await; + + assert_eq!( + cache.matching_pool("team-a", "fingerprint-a").await, + WarmPoolCacheLookup::Match(warm_pool_cache_entry( + "team-a", + "pool-a", + "fingerprint-a" + )) + ); + assert_eq!( + cache.matching_pool("team-a", "fingerprint-b").await, + WarmPoolCacheLookup::NoMatch + ); + }); + } + + #[test] + fn warm_pool_cache_reports_ambiguous_matches() { + futures::executor::block_on(async { + let cache = WarmPoolCache::default(); + cache + .replace_entries(vec![ + warm_pool_cache_entry("team-a", "pool-a", "fingerprint-a"), + warm_pool_cache_entry("team-a", "pool-b", "fingerprint-a"), + ]) + .await; + + assert_eq!( + cache.matching_pool("team-a", "fingerprint-a").await, + WarmPoolCacheLookup::Ambiguous(2) + ); + }); + } + + #[test] + fn template_fingerprint_ignores_sandbox_identity_values() { + let base = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "annotations": { + POD_ANNOTATION_SANDBOX_ID: "sandbox-a", + "stable": "value" + }, + "labels": { + LABEL_SANDBOX_ID: "sandbox-a", + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest", + "env": [ + {"name": openshell_core::sandbox_env::SANDBOX_ID, "value": "sandbox-a"}, + {"name": openshell_core::sandbox_env::SANDBOX, "value": "dev-a"}, + {"name": "ENDPOINT", "value": "https://gateway"} + ] + }] + } + } + } + }); + let changed = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE, + LABEL_SANDBOX_ID: "sandbox-b" + }, + "annotations": { + "stable": "value", + POD_ANNOTATION_SANDBOX_ID: "sandbox-b" + } + }, + "spec": { + "containers": [{ + "env": [ + {"name": openshell_core::sandbox_env::SANDBOX_ID, "value": "sandbox-b"}, + {"name": openshell_core::sandbox_env::SANDBOX, "value": "dev-b"}, + {"name": "ENDPOINT", "value": "https://gateway"} + ], + "image": "example.test/sandbox:latest", + "name": "agent" + }] + } + } + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&base).unwrap(), + sandbox_spec_fingerprint(&changed).unwrap() + ); + } + + #[test] + fn template_fingerprint_ignores_agent_sandbox_defaulted_values() { + let pre_create = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest", + "env": [ + {"name": "OPENSHELL_ENDPOINT", "value": "https://gateway"}, + {"name": openshell_core::sandbox_env::SANDBOX_UID, "value": "1000"} + ] + }], + "initContainers": [{ + "name": "openshell-supervisor-install", + "image": "example.test/supervisor:latest" + }] + } + }, + "volumeClaimTemplates": [{ + "metadata": {"name": "workspace"}, + "spec": {"resources": {"requests": {"storage": "2Gi"}}} + }] + } + }); + let defaulted = serde_json::json!({ + "spec": { + "envVarsInjectionPolicy": "Disallowed", + "networkPolicyManagement": "Managed", + "operatingMode": "Running", + "replicas": 1, + "shutdownPolicy": "Retain", + "volumeClaimTemplatesPolicy": "Disallowed", + "podTemplate": { + "metadata": { + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest", + "env": [ + {"name": "OPENSHELL_ENDPOINT", "value": "https://gateway"}, + {"name": openshell_core::sandbox_env::SANDBOX_UID, "value": "1000"} + ], + "resources": {} + }], + "initContainers": [{ + "name": "openshell-supervisor-install", + "image": "example.test/supervisor:latest", + "resources": {} + }] + } + }, + "volumeClaimTemplates": [{ + "metadata": {"name": "workspace"}, + "spec": {"resources": {"requests": {"storage": "2Gi"}}} + }] + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&pre_create).unwrap(), + sandbox_spec_fingerprint(&defaulted).unwrap() + ); + } + + #[test] + fn template_fingerprint_prunes_empty_identity_metadata() { + let sandbox = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "annotations": { + POD_ANNOTATION_SANDBOX_ID: "sandbox-a" + }, + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE, + LABEL_SANDBOX_ID: "sandbox-a" + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest" + }] + } + } + } + }); + let template = serde_json::json!({ + "spec": { + "podTemplate": { + "metadata": { + "labels": { + LABEL_MANAGED_BY: LABEL_MANAGED_BY_VALUE + } + }, + "spec": { + "containers": [{ + "name": "agent", + "image": "example.test/sandbox:latest" + }] + } + } + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&sandbox).unwrap(), + sandbox_spec_fingerprint(&template).unwrap() + ); + } + + #[test] + fn template_fingerprint_ignores_default_volume_mount_read_only_false() { + let explicit_default = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [{ + "name": "agent", + "image": "image", + "volumeMounts": [{ + "name": "workspace", + "mountPath": "/sandbox", + "readOnly": false + }] + }], + "initContainers": [{ + "name": "openshell-supervisor-install", + "image": "supervisor", + "volumeMounts": [{ + "name": "openshell-supervisor-bin", + "mountPath": "/opt/openshell/bin", + "readOnly": false + }] + }] + } + } + } + }); + let implicit_default = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [{ + "name": "agent", + "image": "image", + "volumeMounts": [{ + "name": "workspace", + "mountPath": "/sandbox" + }] + }], + "initContainers": [{ + "name": "openshell-supervisor-install", + "image": "supervisor", + "volumeMounts": [{ + "name": "openshell-supervisor-bin", + "mountPath": "/opt/openshell/bin" + }] + }] + } + } + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&explicit_default).unwrap(), + sandbox_spec_fingerprint(&implicit_default).unwrap() + ); + } + + #[test] + fn template_fingerprint_ignores_default_cluster_first_dns_policy() { + let explicit_default = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "dnsPolicy": "ClusterFirst", + "containers": [{ + "name": "agent", + "image": "image" + }] + } + } + } + }); + let implicit_default = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [{ + "name": "agent", + "image": "image" + }] + } + } + } + }); + + assert_eq!( + sandbox_spec_fingerprint(&explicit_default).unwrap(), + sandbox_spec_fingerprint(&implicit_default).unwrap() + ); + } + + #[test] + fn sandbox_from_claim_object_preserves_claim_identity() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_CLAIM_KIND); + let mut claim = DynamicObject::new("workspace-a--dev", &resource); + claim.metadata.namespace = Some("team-a".to_string()); + claim.metadata.labels = Some(BTreeMap::from([ + ( + LABEL_MANAGED_BY.to_string(), + LABEL_MANAGED_BY_VALUE.to_string(), + ), + (LABEL_SANDBOX_ID.to_string(), "sandbox-id".to_string()), + (LABEL_SANDBOX_NAME.to_string(), "dev".to_string()), + ( + LABEL_SANDBOX_WORKSPACE.to_string(), + "workspace-a".to_string(), + ), + ])); + claim.data = serde_json::json!({ + "status": { + "sandbox": { + "name": "pool-a-7x9sk" + }, + "conditions": [{ + "type": "Ready", + "status": "False", + "reason": "Binding", + "message": "waiting for warm pool" + }] + } + }); + + let sandbox = sandbox_from_claim_object("default", claim).unwrap(); + + assert_eq!(sandbox.id, "sandbox-id"); + assert_eq!(sandbox.name, "dev"); + assert_eq!(sandbox.workspace, "workspace-a"); + assert_eq!(sandbox.namespace, "team-a"); + assert_eq!( + sandbox + .status + .as_ref() + .map(|status| status.sandbox_name.as_str()), + Some("pool-a-7x9sk") + ); + } + + #[test] + fn sandbox_from_pending_claim_object_has_empty_status() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_CLAIM_KIND); + let mut claim = DynamicObject::new("workspace-a--dev", &resource); + claim.metadata.namespace = Some("team-a".to_string()); + claim.metadata.labels = Some(BTreeMap::from([ + ( + LABEL_MANAGED_BY.to_string(), + LABEL_MANAGED_BY_VALUE.to_string(), + ), + (LABEL_SANDBOX_ID.to_string(), "sandbox-id".to_string()), + (LABEL_SANDBOX_NAME.to_string(), "dev".to_string()), + ( + LABEL_SANDBOX_WORKSPACE.to_string(), + "workspace-a".to_string(), + ), + ])); + + let sandbox = sandbox_from_claim_object("default", claim).unwrap(); + + let status = sandbox.status.expect("pending claim should have status"); + assert!(status.sandbox_name.is_empty()); + assert!(status.instance_id.is_empty()); + assert!(!status.deleting); + } + fn rendered_env<'a>(container: &'a serde_json::Value, name: &str) -> Option<&'a str> { container["env"] .as_array()? @@ -5921,4 +8615,199 @@ mod tests { }; assert!(sandbox_id_from_object(&obj).is_err()); } + + #[test] + fn warm_pool_profile_parses_cli_shaped_schema() { + let profile = parse_warm_pool_profile( + r#" +version = 1 +workspace = "research" +replicas = 3 +image = "ghcr.io/acme/agent-python:latest" +runtime_class_name = "nvidia" + +[environment] +FOO = "bar" + +[resources] +cpu = "4" +memory = "8Gi" +gpu_count = 1 +"#, + ) + .expect("profile should parse"); + + assert_eq!(profile.version, 1); + assert_eq!(profile.workspace, "research"); + assert_eq!(profile.replicas, 3); + assert_eq!(profile.image, "ghcr.io/acme/agent-python:latest"); + assert_eq!(profile.runtime_class_name, "nvidia"); + assert_eq!(profile.environment.get("FOO").unwrap(), "bar"); + assert_eq!(profile.resources.cpu, "4"); + assert_eq!(profile.resources.memory, "8Gi"); + assert_eq!(profile.resources.gpu_count, Some(1)); + } + + #[test] + fn warm_pool_profile_rejects_unknown_fields() { + let err = parse_warm_pool_profile( + r#" +version = 1 +workspace = "research" +replicas = 1 +agent_socket = "/tmp/agent.sock" +"#, + ) + .unwrap_err(); + + assert!(err.contains("unknown field")); + } + + #[test] + fn warm_pool_profile_converts_resources_to_driver_spec() { + let profile = parse_warm_pool_profile( + r#" +version = 1 +workspace = "research" +image = "ghcr.io/acme/agent-python:latest" +runtime_class_name = "kata" + +[environment] +FOO = "bar" + +[resources] +cpu = "500m" +memory = "2Gi" +gpu_count = 2 +"#, + ) + .expect("profile should parse"); + + let spec = warm_pool_profile_to_sandbox_spec(&profile); + assert_eq!(spec.environment.get("FOO").unwrap(), "bar"); + let template = spec.template.as_ref().unwrap(); + assert_eq!(template.image, "ghcr.io/acme/agent-python:latest"); + let resources = template.resources.as_ref().unwrap(); + assert_eq!(resources.cpu_limit, "500m"); + assert_eq!(resources.memory_limit, "2Gi"); + assert_eq!( + platform_config_string(template, "runtime_class_name").as_deref(), + Some("kata") + ); + assert_eq!( + spec.resource_requirements + .as_ref() + .and_then(|requirements| requirements.gpu.as_ref()) + .and_then(|gpu| gpu.count), + Some(2) + ); + } + + #[test] + fn ensure_pod_dns_policy_sets_cluster_first() { + let mut rendered = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [] + } + } + } + }); + + ensure_pod_dns_policy(&mut rendered); + + assert_eq!( + rendered.pointer("/spec/podTemplate/spec/dnsPolicy"), + Some(&serde_json::json!("ClusterFirst")) + ); + } + + #[test] + fn warm_pool_template_identity_env_is_removed() { + let mut rendered = serde_json::json!({ + "spec": { + "podTemplate": { + "spec": { + "containers": [{ + "name": "agent", + "env": [ + {"name": openshell_core::sandbox_env::SANDBOX_ID, "value": ""}, + {"name": openshell_core::sandbox_env::SANDBOX, "value": ""}, + {"name": openshell_core::sandbox_env::ENDPOINT, "value": "http://gateway"} + ] + }] + } + } + } + }); + + remove_warm_pool_template_identity_env(&mut rendered); + + let env = rendered + .pointer("/spec/podTemplate/spec/containers/0/env") + .and_then(serde_json::Value::as_array) + .expect("env should exist"); + let names = env + .iter() + .filter_map(|entry| entry.get("name").and_then(serde_json::Value::as_str)) + .collect::>(); + assert_eq!(names, vec![openshell_core::sandbox_env::ENDPOINT]); + } + + #[test] + fn orphan_gc_only_selects_driver_managed_inactive_profile_resources() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_WARM_POOL_KIND); + let mut active = HashSet::new(); + active.insert("active-uid".to_string()); + + let mut orphan = DynamicObject::new("orphan", &resource); + orphan.metadata.labels = Some(BTreeMap::from([ + ( + LABEL_WARM_POOL_MANAGED_BY.to_string(), + LABEL_WARM_POOL_MANAGED_BY_VALUE.to_string(), + ), + ( + LABEL_WARM_POOL_PROFILE_UID.to_string(), + "deleted-uid".to_string(), + ), + ])); + assert!(warm_pool_profile_generated_resource_is_orphaned( + &orphan, &active + )); + + let mut still_active = DynamicObject::new("active", &resource); + still_active.metadata.labels = Some(BTreeMap::from([ + ( + LABEL_WARM_POOL_MANAGED_BY.to_string(), + LABEL_WARM_POOL_MANAGED_BY_VALUE.to_string(), + ), + ( + LABEL_WARM_POOL_PROFILE_UID.to_string(), + "active-uid".to_string(), + ), + ])); + assert!(!warm_pool_profile_generated_resource_is_orphaned( + &still_active, + &active + )); + + let mut external = DynamicObject::new("external", &resource); + external.metadata.labels = Some(BTreeMap::from([( + LABEL_WARM_POOL_PROFILE_UID.to_string(), + "deleted-uid".to_string(), + )])); + assert!(!warm_pool_profile_generated_resource_is_orphaned( + &external, &active + )); + } + + #[test] + fn generated_warm_pool_profile_name_is_stable_dns_label() { + let name = + generated_warm_pool_profile_name("GPU_Python.Profile", "sha256:abcdef1234567890"); + + assert_eq!(name, "openshell-wp-gpu-python-profile-abcdef12"); + assert!(is_dns_label(&name)); + } } diff --git a/crates/openshell-driver-kubernetes/src/grpc.rs b/crates/openshell-driver-kubernetes/src/grpc.rs index fccfa9464b..5cf8fa6416 100644 --- a/crates/openshell-driver-kubernetes/src/grpc.rs +++ b/crates/openshell-driver-kubernetes/src/grpc.rs @@ -93,11 +93,14 @@ impl ComputeDriver for ComputeDriverService { .into_inner() .sandbox .ok_or_else(|| Status::invalid_argument("sandbox is required"))?; - self.driver + let activation_token = self + .driver .create_sandbox(&sandbox) .await .map_err(|e| Status::from(openshell_core::ComputeDriverError::from(e)))?; - Ok(Response::new(CreateSandboxResponse {})) + Ok(Response::new(CreateSandboxResponse { + activation_token: activation_token.unwrap_or_default(), + })) } async fn stop_sandbox( diff --git a/crates/openshell-driver-kubernetes/src/lib.rs b/crates/openshell-driver-kubernetes/src/lib.rs index 7c56c8de5b..bad24c8759 100644 --- a/crates/openshell-driver-kubernetes/src/lib.rs +++ b/crates/openshell-driver-kubernetes/src/lib.rs @@ -1,14 +1,20 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +pub mod bootstrap; pub mod config; pub mod driver; pub mod grpc; +pub mod sandboxclaim; +pub use bootstrap::{ + K8sIdentityResolver, KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, +}; pub use config::{ AppArmorProfile, DEFAULT_PROXY_UID, DEFAULT_SANDBOX_SERVICE_ACCOUNT_NAME, DEFAULT_WORKSPACE_STORAGE_SIZE, KubernetesComputeConfig, KubernetesSidecarConfig, - SupervisorSideloadMethod, SupervisorTopology, + KubernetesWarmPoolingConfig, SupervisorSideloadMethod, SupervisorTopology, }; pub use driver::{KubernetesComputeDriver, KubernetesDriverError}; pub use grpc::ComputeDriverService; +pub use sandboxclaim::SandboxClaimActivationController; diff --git a/crates/openshell-driver-kubernetes/src/main.rs b/crates/openshell-driver-kubernetes/src/main.rs index c733b8a45b..14e23dc0b0 100644 --- a/crates/openshell-driver-kubernetes/src/main.rs +++ b/crates/openshell-driver-kubernetes/src/main.rs @@ -12,7 +12,7 @@ use openshell_core::proto::compute::v1::compute_driver_server::ComputeDriverServ use openshell_driver_kubernetes::{ AppArmorProfile, ComputeDriverService, DEFAULT_PROXY_UID, DEFAULT_SANDBOX_SERVICE_ACCOUNT_NAME, KubernetesComputeConfig, KubernetesComputeDriver, KubernetesSidecarConfig, - SupervisorSideloadMethod, SupervisorTopology, + KubernetesWarmPoolingConfig, SupervisorSideloadMethod, SupervisorTopology, }; #[derive(Parser, Debug)] @@ -107,8 +107,8 @@ struct Args { app_armor_profile: Option, /// Lifetime (seconds) of the projected `ServiceAccount` token - /// kubelet writes into each sandbox pod for the `IssueSandboxToken` - /// bootstrap exchange. Kubelet enforces a minimum of 600s; the + /// kubelet writes into each sandbox pod for the `RegisterSupervisorPod` + /// bootstrap stream. Kubelet enforces a minimum of 600s; the /// gateway clamps values outside `[600, 86400]`. Default 3600. #[arg(long, env = "OPENSHELL_K8S_SA_TOKEN_TTL_SECS", default_value_t = 3600)] sa_token_ttl_secs: i64, @@ -148,6 +148,7 @@ async fn main() -> Result<()> { proxy_uid: args.sidecar_proxy_uid, process_binary_aware_network_policy: args.sidecar_process_binary_aware_network_policy, }, + warm_pooling: KubernetesWarmPoolingConfig::default(), grpc_endpoint: args.grpc_endpoint.unwrap_or_default(), ssh_socket_path: args.sandbox_ssh_socket_path, client_tls_secret_name: args.client_tls_secret_name.unwrap_or_default(), diff --git a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs new file mode 100644 index 0000000000..d1f361f5d8 --- /dev/null +++ b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs @@ -0,0 +1,1042 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Kubernetes `SandboxClaim` activation support. + +use base64::Engine as _; +use futures::{StreamExt, TryStreamExt}; +use k8s_openapi::api::core::v1::Pod; +use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference; +use kube::api::{Api, ApiResource, ListParams}; +use kube::core::DynamicObject; +use kube::core::gvk::GroupVersionKind; +use kube::runtime::watcher::{self, Event}; +use kube::{Client, Error as KubeError}; +use openshell_core::driver_utils::{ + LABEL_MANAGED_BY, LABEL_MANAGED_BY_VALUE, LABEL_SANDBOX_ID, LABEL_SANDBOX_WORKSPACE, +}; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapActivationRequest, SupervisorBootstrapActivator, +}; +use sha2::{Digest, Sha256}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::watch; +use tonic::Code; +use tracing::{debug, info, warn}; + +const KUBE_API_TIMEOUT: Duration = Duration::from_secs(30); +const ACTIVATION_RETRY_ATTEMPTS: usize = 8; +const ACTIVATION_RETRY_INITIAL_DELAY: Duration = Duration::from_millis(250); +const ACTIVATION_RETRY_MAX_DELAY: Duration = Duration::from_secs(5); +const DRIVER_NAME: &str = "kubernetes"; + +const SANDBOX_GROUP: &str = "agents.x-k8s.io"; +const SANDBOX_KIND: &str = "Sandbox"; +const SANDBOX_VERSION_V1BETA1: &str = "v1beta1"; +const SANDBOX_VERSION_V1ALPHA1: &str = "v1alpha1"; +const SANDBOX_VERSIONS: &[&str] = &[SANDBOX_VERSION_V1BETA1, SANDBOX_VERSION_V1ALPHA1]; +const SANDBOX_API_VERSION_FULL_V1BETA1: &str = "agents.x-k8s.io/v1beta1"; +const SANDBOX_API_VERSION_FULL_V1ALPHA1: &str = "agents.x-k8s.io/v1alpha1"; +const SANDBOX_POD_NAME_ANNOTATION: &str = "agents.x-k8s.io/pod-name"; + +const SANDBOX_CLAIM_GROUP: &str = "extensions.agents.x-k8s.io"; +const SANDBOX_CLAIM_KIND: &str = "SandboxClaim"; +const SANDBOX_CLAIM_VERSION_V1BETA1: &str = "v1beta1"; +const LABEL_ALLOCATION: &str = "openshell.ai/allocation"; +const LABEL_ALLOCATION_SANDBOX_CLAIM: &str = "sandbox-claim"; + +/// Watches Kubernetes `SandboxClaim` resources and activates pending +/// supervisor bootstrap streams after revalidating live driver state. +#[derive(Clone)] +pub struct SandboxClaimActivationController { + client: Client, + watch_client: Client, + namespace: String, +} + +impl SandboxClaimActivationController { + #[must_use] + pub fn new(client: Client, watch_client: Client, namespace: impl Into) -> Self { + Self { + client, + watch_client, + namespace: namespace.into(), + } + } + + pub fn spawn( + &self, + activator: Arc, + shutdown_rx: watch::Receiver, + ) { + let controller = self.clone(); + tokio::spawn(async move { + controller.run(activator, shutdown_rx).await; + }); + } + + async fn run( + self, + activator: Arc, + mut shutdown_rx: watch::Receiver, + ) { + let claim_api = match self + .supported_sandbox_claim_api(self.watch_client.clone()) + .await + { + Ok(api) => api, + Err(err) => { + debug!( + namespace = %self.namespace, + error = %err, + "SandboxClaim API is not available; warm-pool claim activation disabled" + ); + return; + } + }; + + let mut stream = watcher::watcher(claim_api.api, watcher::Config::default()).boxed(); + info!( + namespace = %self.namespace, + sandbox_claim_api_version = %claim_api.version, + "Watching Kubernetes SandboxClaims for warm-pool activation" + ); + + loop { + tokio::select! { + changed = shutdown_rx.changed() => { + if changed.is_err() || *shutdown_rx.borrow() { + break; + } + } + event = stream.try_next() => match event { + Ok(Some(Event::Applied(claim))) => { + self.handle_claim(claim, activator.as_ref()).await; + } + Ok(Some(Event::Restarted(claims))) => { + for claim in claims { + self.handle_claim(claim, activator.as_ref()).await; + } + } + Ok(Some(Event::Deleted(_))) => {} + Ok(None) => break, + Err(err) => { + warn!( + namespace = %self.namespace, + error = %err, + "SandboxClaim watch failed; warm-pool claim activation stopped" + ); + break; + } + } + } + } + } + + async fn handle_claim( + &self, + claim: DynamicObject, + activator: &(dyn SupervisorBootstrapActivator + '_), + ) { + let claim = match parse_sandbox_claim(&claim) { + Ok(claim) => claim, + Err(err) => { + warn!( + namespace = %self.namespace, + error = %err, + "Ignoring invalid SandboxClaim" + ); + return; + } + }; + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox_claim_uid = %claim.uid, + warm_pool = ?claim.warm_pool_name, + sandbox_template = ?claim.sandbox_template_name, + sandbox = ?claim.sandbox_name, + pod_ips = ?claim.pod_ips, + "Processing SandboxClaim for warm-pool activation" + ); + let Some(sandbox_name) = claim.sandbox_name.as_deref() else { + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + "SandboxClaim has not selected a Sandbox yet" + ); + return; + }; + + let sandbox_cr = match self.get_sandbox_cr(sandbox_name).await { + Ok(Some(sandbox_cr)) => sandbox_cr, + Ok(None) => { + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + "SandboxClaim selected Sandbox is not available yet" + ); + return; + } + Err(err) => { + warn!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + error = %err, + "Failed to read Sandbox selected by SandboxClaim" + ); + return; + } + }; + + let pod = match self.resolve_controlled_pod(&sandbox_cr).await { + Ok(Some(pod)) => pod, + Ok(None) => return, + Err(err) => { + warn!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + error = %err, + "Failed to resolve Sandbox pod for SandboxClaim activation" + ); + return; + } + }; + + let request = match activation_request_from_claim_state(&claim, &sandbox_cr, &pod) { + Ok(Some(request)) => request, + Ok(None) => return, + Err(err) => { + warn!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + error = %err, + "SandboxClaim activation validation failed" + ); + return; + } + }; + + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + owner_uid = %request.owner_uid, + pod = %pod.metadata.name.as_deref().unwrap_or_default(), + "Resolved SandboxClaim activation target" + ); + + activate_registered_supervisor_with_retry( + activator, + request, + ActivationLogContext { + namespace: self.namespace.as_str(), + sandbox_claim: claim.name.as_str(), + sandbox: sandbox_name, + }, + ACTIVATION_RETRY_ATTEMPTS, + ACTIVATION_RETRY_INITIAL_DELAY, + ACTIVATION_RETRY_MAX_DELAY, + ) + .await; + } + + fn sandbox_claim_api(&self, client: Client, version: &'static str) -> SandboxClaimApi { + let gvk = GroupVersionKind::gvk(SANDBOX_CLAIM_GROUP, version, SANDBOX_CLAIM_KIND); + let resource = ApiResource::from_gvk(&gvk); + let api = Api::namespaced_with(client, &self.namespace, &resource); + SandboxClaimApi { api, version } + } + + async fn supported_sandbox_claim_api(&self, client: Client) -> Result { + let claim_api = self.sandbox_claim_api(client, SANDBOX_CLAIM_VERSION_V1BETA1); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + claim_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => Ok(claim_api), + Ok(Err(err)) if should_try_next_api_version(&err) => Err( + "SandboxClaim API extensions.agents.x-k8s.io/v1beta1 is not available".to_string(), + ), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } + } + + fn sandbox_api(&self, client: Client, version: &'static str) -> Api { + let gvk = GroupVersionKind::gvk(SANDBOX_GROUP, version, SANDBOX_KIND); + let resource = ApiResource::from_gvk(&gvk); + Api::namespaced_with(client, &self.namespace, &resource) + } + + async fn get_sandbox_cr(&self, name: &str) -> Result, String> { + for version in SANDBOX_VERSIONS { + let sandbox_api = self.sandbox_api(self.client.clone(), version); + match tokio::time::timeout(KUBE_API_TIMEOUT, sandbox_api.get(name)).await { + Ok(Ok(sandbox_cr)) => return Ok(Some(sandbox_cr)), + Ok(Err(KubeError::Api(err))) if err.code == 404 => {} + Ok(Err(err)) if should_try_next_api_version(&err) => {} + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + } + } + Ok(None) + } + + async fn resolve_controlled_pod( + &self, + sandbox_cr: &DynamicObject, + ) -> Result, String> { + let Some(owner_uid) = sandbox_cr.metadata.uid.as_deref() else { + return Err("Sandbox CR is missing uid".to_string()); + }; + + let pods_api: Api = Api::namespaced(self.client.clone(), &self.namespace); + if let Some(pod_name) = sandbox_pod_name(sandbox_cr) { + let pod = match tokio::time::timeout(KUBE_API_TIMEOUT, pods_api.get(&pod_name)).await { + Ok(Ok(pod)) => pod, + Ok(Err(KubeError::Api(err))) if err.code == 404 => { + debug!( + namespace = %self.namespace, + sandbox = %sandbox_cr.metadata.name.as_deref().unwrap_or_default(), + sandbox_uid = %owner_uid, + pod = %pod_name, + "Annotated Sandbox pod was not found for SandboxClaim activation" + ); + return Ok(None); + } + Ok(Err(err)) => return Err(err.to_string()), + Err(_) => { + return Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )); + } + }; + if !pod_has_sandbox_owner(&pod, owner_uid) { + return Err(format!( + "annotated Sandbox pod {pod_name} is not controlled by the selected Sandbox" + )); + } + return Ok(Some(pod)); + } + + let pods = match sandbox_selector(sandbox_cr) { + Some(selector) => { + let params = ListParams::default().labels(&selector); + list_pods(&pods_api, ¶ms).await? + } + None => list_pods(&pods_api, &ListParams::default()).await?, + }; + + let controlled = pods + .into_iter() + .filter(|pod| pod_has_sandbox_owner(pod, owner_uid)) + .collect::>(); + + match controlled.as_slice() { + [pod] => Ok(Some(pod.clone())), + [] => { + debug!( + namespace = %self.namespace, + sandbox = %sandbox_cr.metadata.name.as_deref().unwrap_or_default(), + sandbox_uid = %owner_uid, + "No controlled pod found for SandboxClaim activation" + ); + Ok(None) + } + _ => Err(format!( + "expected one controlled Sandbox pod, found {}", + controlled.len() + )), + } + } +} + +struct SandboxClaimApi { + api: Api, + version: &'static str, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +struct ParsedSandboxClaim { + name: String, + namespace: String, + uid: String, + sandbox_id: Option, + workspace: Option, + managed_by: Option, + allocation: Option, + warm_pool_name: Option, + sandbox_template_name: Option, + sandbox_name: Option, + pod_ips: Vec, +} + +struct ActivationLogContext<'a> { + namespace: &'a str, + sandbox_claim: &'a str, + sandbox: &'a str, +} + +async fn activate_registered_supervisor_with_retry( + activator: &(dyn SupervisorBootstrapActivator + '_), + request: SupervisorBootstrapActivationRequest, + context: ActivationLogContext<'_>, + attempts: usize, + initial_delay: Duration, + max_delay: Duration, +) { + let attempts = attempts.max(1); + let mut delay = initial_delay; + + for attempt in 1..=attempts { + match activator + .activate_registered_supervisor(request.clone()) + .await + { + Ok(()) => { + info!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempt, + "Activated warm-pool supervisor from SandboxClaim" + ); + return; + } + Err(status) if status.code() == Code::AlreadyExists => { + debug!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempt, + "Warm-pool supervisor was already activated" + ); + return; + } + Err(status) if status.code() == Code::NotFound && attempt < attempts => { + debug!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempt, + retry_after_ms = delay.as_millis(), + "Warm-pool supervisor registration is not pending yet; retrying activation" + ); + tokio::time::sleep(delay).await; + delay = (delay * 2).min(max_delay); + } + Err(status) if status.code() == Code::NotFound => { + debug!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempts, + "Warm-pool supervisor registration is not pending after retries" + ); + return; + } + Err(status) => { + warn!( + namespace = %context.namespace, + sandbox_claim = %context.sandbox_claim, + sandbox = %context.sandbox, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + attempt, + code = ?status.code(), + message = %status.message(), + "Failed to activate warm-pool supervisor from SandboxClaim" + ); + return; + } + } + } +} + +fn parse_sandbox_claim(obj: &DynamicObject) -> Result { + Ok(ParsedSandboxClaim { + name: required_metadata_field(&obj.metadata.name, "name")?, + namespace: obj.metadata.namespace.clone().unwrap_or_default(), + uid: required_metadata_field(&obj.metadata.uid, "uid")?, + sandbox_id: label_value(obj, LABEL_SANDBOX_ID), + workspace: label_value(obj, LABEL_SANDBOX_WORKSPACE), + managed_by: label_value(obj, LABEL_MANAGED_BY), + allocation: label_value(obj, LABEL_ALLOCATION), + warm_pool_name: string_at(&obj.data, &["spec", "warmPoolRef", "name"]), + sandbox_template_name: string_at(&obj.data, &["spec", "sandboxTemplateRef", "name"]), + sandbox_name: string_at(&obj.data, &["status", "sandbox", "name"]), + pod_ips: strings_at(&obj.data, &["status", "sandbox", "podIPs"]), + }) +} + +fn label_value(obj: &DynamicObject, key: &str) -> Option { + obj.metadata + .labels + .as_ref() + .and_then(|labels| labels.get(key)) + .filter(|value| !value.is_empty()) + .cloned() +} + +fn required_metadata_field(value: &Option, field: &str) -> Result { + value + .as_ref() + .filter(|value| !value.is_empty()) + .cloned() + .ok_or_else(|| format!("SandboxClaim is missing metadata.{field}")) +} + +pub(crate) fn sandbox_claim_activation_token(namespace: &str, name: &str, uid: &str) -> String { + let input = + format!("openshell-kubernetes-sandbox-claim-activation-v1\0{namespace}\0{name}\0{uid}"); + let digest = Sha256::digest(input.as_bytes()); + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest) +} + +fn activation_request_from_claim_state( + claim: &ParsedSandboxClaim, + sandbox_cr: &DynamicObject, + pod: &Pod, +) -> Result, String> { + let Some(claim_sandbox_name) = claim.sandbox_name.as_deref() else { + return Ok(None); + }; + let sandbox_name = sandbox_cr.metadata.name.as_deref().unwrap_or_default(); + if sandbox_name != claim_sandbox_name { + return Err(format!( + "SandboxClaim selected Sandbox {claim_sandbox_name}, but live Sandbox is {sandbox_name}" + )); + } + if claim.managed_by.as_deref() != Some(LABEL_MANAGED_BY_VALUE) { + return Err("SandboxClaim is not OpenShell-managed".to_string()); + } + if claim.allocation.as_deref() != Some(LABEL_ALLOCATION_SANDBOX_CLAIM) { + return Err("SandboxClaim is missing OpenShell sandbox-claim allocation label".to_string()); + } + let owner_uid = sandbox_cr + .metadata + .uid + .as_ref() + .filter(|uid| !uid.is_empty()) + .cloned() + .ok_or_else(|| "Sandbox CR is missing uid".to_string())?; + let sandbox_id = claim + .sandbox_id + .clone() + .ok_or_else(|| "SandboxClaim is missing OpenShell sandbox id label".to_string())?; + if let Some(live_sandbox_id) = sandbox_cr + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(LABEL_SANDBOX_ID)) + .filter(|id| !id.is_empty()) + && live_sandbox_id != &sandbox_id + { + return Err("SandboxClaim sandbox id does not match selected Sandbox label".to_string()); + } + if let Some(claim_workspace) = claim.workspace.as_deref() + && let Some(sandbox_workspace) = sandbox_cr + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(LABEL_SANDBOX_WORKSPACE)) + .filter(|workspace| !workspace.is_empty()) + && sandbox_workspace != claim_workspace + { + return Err("SandboxClaim workspace does not match selected Sandbox label".to_string()); + } + let instance_id = pod + .metadata + .uid + .as_ref() + .filter(|uid| !uid.is_empty()) + .cloned() + .ok_or_else(|| "Sandbox pod is missing uid".to_string())?; + if !pod_has_sandbox_owner(pod, &owner_uid) { + return Err("Sandbox pod is not controlled by the selected Sandbox".to_string()); + } + + Ok(Some(SupervisorBootstrapActivationRequest { + driver: DRIVER_NAME.to_string(), + instance_id, + sandbox_id, + owner_uid, + activation_token: sandbox_claim_activation_token(&claim.namespace, &claim.name, &claim.uid), + reason: format!("SandboxClaim/{}", claim.name), + })) +} + +async fn list_pods(api: &Api, params: &ListParams) -> Result, String> { + match tokio::time::timeout(KUBE_API_TIMEOUT, api.list(params)).await { + Ok(Ok(list)) => Ok(list.items), + Ok(Err(err)) => Err(err.to_string()), + Err(_) => Err(format!( + "timed out after {}s waiting for Kubernetes API", + KUBE_API_TIMEOUT.as_secs() + )), + } +} + +fn sandbox_selector(obj: &DynamicObject) -> Option { + string_at(&obj.data, &["status", "selector"]) +} + +fn sandbox_pod_name(obj: &DynamicObject) -> Option { + obj.metadata + .annotations + .as_ref() + .and_then(|annotations| annotations.get(SANDBOX_POD_NAME_ANNOTATION)) + .filter(|pod_name| !pod_name.is_empty()) + .cloned() +} + +fn pod_has_sandbox_owner(pod: &Pod, owner_uid: &str) -> bool { + pod.metadata + .owner_references + .as_deref() + .unwrap_or_default() + .iter() + .any(|owner| is_supported_sandbox_owner_reference(owner) && owner.uid == owner_uid) +} + +fn is_supported_sandbox_owner_reference(owner: &OwnerReference) -> bool { + owner.kind == SANDBOX_KIND + && owner.controller == Some(true) + && matches!( + owner.api_version.as_str(), + SANDBOX_API_VERSION_FULL_V1BETA1 | SANDBOX_API_VERSION_FULL_V1ALPHA1 + ) +} + +fn should_try_next_api_version(err: &KubeError) -> bool { + matches!(err, KubeError::Api(api) if api.code == 404) +} + +fn string_at(value: &serde_json::Value, path: &[&str]) -> Option { + path.iter() + .try_fold(value, |current, segment| current.get(segment)) + .and_then(|value| value.as_str()) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn strings_at(value: &serde_json::Value, path: &[&str]) -> Vec { + path.iter() + .try_fold(value, |current, segment| current.get(segment)) + .and_then(|value| value.as_array()) + .into_iter() + .flatten() + .filter_map(|value| value.as_str()) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use kube::core::ObjectMeta; + use serde_json::json; + use std::collections::BTreeMap; + use std::sync::Mutex; + use tonic::Status; + use tonic::async_trait; + + struct FakeActivator { + outcomes: Mutex>>, + requests: Mutex>, + } + + impl FakeActivator { + fn new(outcomes: Vec>) -> Self { + Self { + outcomes: Mutex::new(outcomes), + requests: Mutex::new(Vec::new()), + } + } + + fn request_count(&self) -> usize { + self.requests.lock().expect("requests mutex poisoned").len() + } + } + + #[async_trait] + impl SupervisorBootstrapActivator for FakeActivator { + async fn activate_registered_supervisor( + &self, + request: SupervisorBootstrapActivationRequest, + ) -> Result<(), Status> { + self.requests + .lock() + .expect("requests mutex poisoned") + .push(request); + self.outcomes + .lock() + .expect("outcomes mutex poisoned") + .remove(0) + } + } + + fn dynamic_object( + group: &'static str, + version: &'static str, + kind: &'static str, + name: &str, + uid: &str, + data: serde_json::Value, + ) -> DynamicObject { + let gvk = GroupVersionKind::gvk(group, version, kind); + let resource = ApiResource::from_gvk(&gvk); + let mut obj = DynamicObject::new(name, &resource); + obj.metadata.namespace = Some("openshell".to_string()); + obj.metadata.uid = Some(uid.to_string()); + obj.data = data; + obj + } + + fn claim_labels(sandbox_id: &str) -> BTreeMap { + BTreeMap::from([ + ( + LABEL_MANAGED_BY.to_string(), + LABEL_MANAGED_BY_VALUE.to_string(), + ), + ( + LABEL_ALLOCATION.to_string(), + LABEL_ALLOCATION_SANDBOX_CLAIM.to_string(), + ), + (LABEL_SANDBOX_ID.to_string(), sandbox_id.to_string()), + (LABEL_SANDBOX_WORKSPACE.to_string(), "default".to_string()), + ]) + } + + fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str, version: &'static str) -> DynamicObject { + let mut obj = dynamic_object( + SANDBOX_GROUP, + version, + SANDBOX_KIND, + name, + uid, + json!({"status": {"selector": "agents.x-k8s.io/sandbox=sandbox-a"}}), + ); + obj.metadata + .labels + .get_or_insert_with(BTreeMap::new) + .insert(LABEL_SANDBOX_ID.to_string(), sandbox_id.to_string()); + obj.metadata + .labels + .get_or_insert_with(BTreeMap::new) + .insert(LABEL_SANDBOX_WORKSPACE.to_string(), "default".to_string()); + obj + } + + fn sandbox_pod(name: &str, uid: &str, owner_uid: &str, api_version: &str) -> Pod { + Pod { + metadata: ObjectMeta { + name: Some(name.to_string()), + uid: Some(uid.to_string()), + owner_references: Some(vec![OwnerReference { + api_version: api_version.to_string(), + kind: SANDBOX_KIND.to_string(), + name: "sandbox-a".to_string(), + uid: owner_uid.to_string(), + controller: Some(true), + block_owner_deletion: None, + }]), + ..ObjectMeta::default() + }, + ..Pod::default() + } + } + + #[test] + fn parses_v1beta1_sandbox_claim_status() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({ + "spec": {"warmPoolRef": {"name": "pool-a"}}, + "status": { + "sandbox": { + "name": "sandbox-a", + "podIPs": ["10.0.0.10", "fd00::10"] + } + } + }), + ); + + let claim = parse_sandbox_claim(&obj).unwrap(); + + assert_eq!(claim.name, "claim-a"); + assert_eq!(claim.uid, "claim-uid"); + assert_eq!(claim.sandbox_id, None); + assert_eq!(claim.warm_pool_name.as_deref(), Some("pool-a")); + assert_eq!(claim.sandbox_name.as_deref(), Some("sandbox-a")); + assert_eq!(claim.pod_ips, vec!["10.0.0.10", "fd00::10"]); + } + + #[test] + fn claim_without_selected_sandbox_does_not_activate() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"spec": {"warmPoolRef": {"name": "pool-a"}}}), + ); + let claim = parse_sandbox_claim(&obj).unwrap(); + let sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1BETA1, + ); + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "sandbox-uid", + SANDBOX_API_VERSION_FULL_V1BETA1, + ); + + assert_eq!( + activation_request_from_claim_state(&claim, &sandbox, &pod).unwrap(), + None + ); + } + + #[test] + fn activation_request_uses_live_sandbox_and_pod_identity() { + let mut obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({ + "spec": {"warmPoolRef": {"name": "pool-a"}}, + "status": {"sandbox": {"name": "sandbox-a"}} + }), + ); + obj.metadata.labels = Some(claim_labels("openshell-sandbox")); + let claim = parse_sandbox_claim(&obj).unwrap(); + let sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1ALPHA1, + ); + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "sandbox-uid", + SANDBOX_API_VERSION_FULL_V1ALPHA1, + ); + + let request = activation_request_from_claim_state(&claim, &sandbox, &pod) + .unwrap() + .unwrap(); + + assert_eq!(request.driver, "kubernetes"); + assert_eq!(request.instance_id, "pod-uid"); + assert_eq!(request.sandbox_id, "openshell-sandbox"); + assert_eq!(request.owner_uid, "sandbox-uid"); + assert_eq!( + request.activation_token, + sandbox_claim_activation_token("openshell", "claim-a", "claim-uid") + ); + assert_eq!(request.reason, "SandboxClaim/claim-a"); + } + + #[test] + fn activation_request_uses_claim_sandbox_id_when_sandbox_lacks_label() { + let mut obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"status": {"sandbox": {"name": "sandbox-a"}}}), + ); + obj.metadata.labels = Some(claim_labels("openshell-sandbox")); + let claim = parse_sandbox_claim(&obj).unwrap(); + let mut sandbox = sandbox_cr("sandbox-a", "sandbox-uid", "", SANDBOX_VERSION_V1BETA1); + sandbox.metadata.labels = None; + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "sandbox-uid", + SANDBOX_API_VERSION_FULL_V1BETA1, + ); + + let request = activation_request_from_claim_state(&claim, &sandbox, &pod) + .unwrap() + .unwrap(); + + assert_eq!(request.sandbox_id, "openshell-sandbox"); + } + + #[tokio::test] + async fn activation_retry_handles_registration_after_claim_binding() { + let activator = FakeActivator::new(vec![Err(Status::not_found("not pending")), Ok(())]); + let request = SupervisorBootstrapActivationRequest { + driver: DRIVER_NAME.to_string(), + instance_id: "pod-uid".to_string(), + sandbox_id: "sandbox-id".to_string(), + owner_uid: "sandbox-uid".to_string(), + activation_token: "activation-token".to_string(), + reason: "SandboxClaim/claim-a".to_string(), + }; + + activate_registered_supervisor_with_retry( + &activator, + request, + ActivationLogContext { + namespace: "openshell", + sandbox_claim: "claim-a", + sandbox: "sandbox-a", + }, + 2, + Duration::ZERO, + Duration::ZERO, + ) + .await; + + assert_eq!(activator.request_count(), 2); + } + + #[test] + fn activation_request_rejects_missing_sandbox_id_labels() { + let mut obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"status": {"sandbox": {"name": "sandbox-a"}}}), + ); + let mut labels = claim_labels(""); + labels.remove(LABEL_SANDBOX_ID); + obj.metadata.labels = Some(labels); + let claim = parse_sandbox_claim(&obj).unwrap(); + let mut sandbox = sandbox_cr("sandbox-a", "sandbox-uid", "", SANDBOX_VERSION_V1BETA1); + sandbox.metadata.labels = None; + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "sandbox-uid", + SANDBOX_API_VERSION_FULL_V1BETA1, + ); + + let err = activation_request_from_claim_state(&claim, &sandbox, &pod).unwrap_err(); + + assert!(err.contains("missing OpenShell sandbox id label")); + } + + #[test] + fn activation_request_rejects_uncontrolled_pod() { + let mut obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"status": {"sandbox": {"name": "sandbox-a"}}}), + ); + obj.metadata.labels = Some(claim_labels("openshell-sandbox")); + let claim = parse_sandbox_claim(&obj).unwrap(); + let sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1BETA1, + ); + let pod = sandbox_pod( + "pod-a", + "pod-uid", + "other-sandbox-uid", + SANDBOX_API_VERSION_FULL_V1BETA1, + ); + + assert!(activation_request_from_claim_state(&claim, &sandbox, &pod).is_err()); + } + + #[test] + fn sandbox_selector_reads_status_selector() { + let sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1BETA1, + ); + + assert_eq!( + sandbox_selector(&sandbox).as_deref(), + Some("agents.x-k8s.io/sandbox=sandbox-a") + ); + } + + #[test] + fn sandbox_pod_name_reads_pod_name_annotation() { + let mut sandbox = sandbox_cr( + "sandbox-a", + "sandbox-uid", + "openshell-sandbox", + SANDBOX_VERSION_V1BETA1, + ); + sandbox.metadata.annotations = Some(BTreeMap::from([( + SANDBOX_POD_NAME_ANNOTATION.to_string(), + "pod-a".to_string(), + )])); + + assert_eq!(sandbox_pod_name(&sandbox).as_deref(), Some("pod-a")); + } + + #[test] + fn api_version_probe_retries_only_404_errors() { + let unavailable = KubeError::Api(kube::core::ErrorResponse { + status: "404 Not Found".to_string(), + message: "could not find the requested resource".to_string(), + reason: "NotFound".to_string(), + code: 404, + }); + let forbidden = KubeError::Api(kube::core::ErrorResponse { + status: "Failure".to_string(), + message: "forbidden".to_string(), + reason: "Forbidden".to_string(), + code: 403, + }); + + assert!(should_try_next_api_version(&unavailable)); + assert!(!should_try_next_api_version(&forbidden)); + } +} diff --git a/crates/openshell-driver-podman/src/grpc.rs b/crates/openshell-driver-podman/src/grpc.rs index 8e68a91e72..10cd85712b 100644 --- a/crates/openshell-driver-podman/src/grpc.rs +++ b/crates/openshell-driver-podman/src/grpc.rs @@ -96,7 +96,9 @@ impl ComputeDriver for ComputeDriverService { .create_sandbox(&sandbox) .await .map_err(Status::from)?; - Ok(Response::new(CreateSandboxResponse {})) + Ok(Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( diff --git a/crates/openshell-driver-vm/src/driver.rs b/crates/openshell-driver-vm/src/driver.rs index 7af0ddc389..ab6e9c99cd 100644 --- a/crates/openshell-driver-vm/src/driver.rs +++ b/crates/openshell-driver-vm/src/driver.rs @@ -622,7 +622,9 @@ impl VmDriver { task.abort(); } - Ok(CreateSandboxResponse {}) + Ok(CreateSandboxResponse { + activation_token: String::new(), + }) } async fn provision_sandbox( diff --git a/crates/openshell-sandbox/src/lib.rs b/crates/openshell-sandbox/src/lib.rs index 271f82710c..0b824d73c9 100644 --- a/crates/openshell-sandbox/src/lib.rs +++ b/crates/openshell-sandbox/src/lib.rs @@ -102,6 +102,8 @@ pub async fn run_sandbox( let (program, args) = command .split_first() .ok_or_else(|| miette::miette!("No command specified"))?; + let mut sandbox_id = sandbox_id; + let mut sandbox = sandbox; // Initialize the process-wide OCSF context early so that events emitted // during policy loading (filesystem config, validation) have a context. @@ -151,6 +153,19 @@ pub async fn run_sandbox( // Load policy and initialize OPA engine let openshell_endpoint_for_proxy = openshell_endpoint.clone(); + if sidecar_bootstrap.is_none() + && sandbox_id.is_none() + && let Some(endpoint) = openshell_endpoint.as_deref() + && std::env::var(openshell_core::sandbox_env::K8S_SA_TOKEN_FILE) + .ok() + .is_some_and(|path| !path.is_empty()) + { + let activation = openshell_core::grpc_client::register_k8s_supervisor_pod(endpoint).await?; + if sandbox.is_none() && !activation.sandbox_name.is_empty() { + sandbox = Some(activation.sandbox_name); + } + sandbox_id = Some(activation.sandbox_id); + } let sandbox_name_for_agg = sandbox.clone(); let ( mut policy, @@ -2085,7 +2100,8 @@ async fn load_policy( Err(miette::miette!( "Sandbox policy required. Provide one of:\n\ - --policy-rules and --policy-data (or OPENSHELL_POLICY_RULES and OPENSHELL_POLICY_DATA env vars)\n\ - - --sandbox-id and --openshell-endpoint (or OPENSHELL_SANDBOX_ID and OPENSHELL_ENDPOINT env vars)" + - --sandbox-id and --openshell-endpoint (or OPENSHELL_SANDBOX_ID and OPENSHELL_ENDPOINT env vars)\n\ + - OPENSHELL_ENDPOINT and OPENSHELL_K8S_SA_TOKEN_FILE for Kubernetes supervisor registration" )) } diff --git a/crates/openshell-sdk/tests/client_mock.rs b/crates/openshell-sdk/tests/client_mock.rs index d0512d26e1..5b24032b17 100644 --- a/crates/openshell-sdk/tests/client_mock.rs +++ b/crates/openshell-sdk/tests/client_mock.rs @@ -243,6 +243,16 @@ impl OpenShell for TestOpenShell { Ok(Response::new(proto::CreateSshSessionResponse::default())) } + type RegisterSupervisorPodStream = + tokio_stream::wrappers::ReceiverStream>; + + async fn register_supervisor_pod( + &self, + _: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("unused")) + } + async fn expose_service( &self, _: tonic::Request, diff --git a/crates/openshell-server-macros/src/lib.rs b/crates/openshell-server-macros/src/lib.rs index a698ae6623..d085080f48 100644 --- a/crates/openshell-server-macros/src/lib.rs +++ b/crates/openshell-server-macros/src/lib.rs @@ -54,6 +54,7 @@ struct RpcAuth { enum AuthMode { Unauthenticated, Sandbox, + PodRegistration, Bearer, Dual, } @@ -116,11 +117,11 @@ impl RpcAuth { }; match mode { - AuthMode::Unauthenticated | AuthMode::Sandbox => { + AuthMode::Unauthenticated | AuthMode::Sandbox | AuthMode::PodRegistration => { if let Some(ref s) = scope { return Err(Error::new( s.span(), - "`scope` is only valid for `auth = \"bearer\"` or `auth = \"dual\"` (sandbox principals don't carry scopes)", + "`scope` is only valid for `auth = \"bearer\"` or `auth = \"dual\"` (sandbox and pod-registration principals don't carry scopes)", )); } if role.is_some() { @@ -154,12 +155,13 @@ fn parse_auth_mode(value: &LitStr) -> Result { match value.value().as_str() { "unauthenticated" => Ok(AuthMode::Unauthenticated), "sandbox" => Ok(AuthMode::Sandbox), + "pod_registration" => Ok(AuthMode::PodRegistration), "bearer" => Ok(AuthMode::Bearer), "dual" => Ok(AuthMode::Dual), other => Err(Error::new( value.span(), format!( - "invalid auth mode `{other}`; expected one of `unauthenticated`, `sandbox`, `bearer`, `dual`" + "invalid auth mode `{other}`; expected one of `unauthenticated`, `sandbox`, `pod_registration`, `bearer`, `dual`" ), )), } @@ -288,6 +290,9 @@ fn expand(args: &AuthzArgs, item: &mut ItemImpl) -> Result { quote! { crate::auth::method_authz::AuthMode::Sandbox } } + AuthMode::PodRegistration => { + quote! { crate::auth::method_authz::AuthMode::PodRegistration } + } AuthMode::Bearer => quote! { crate::auth::method_authz::AuthMode::Bearer }, AuthMode::Dual => quote! { crate::auth::method_authz::AuthMode::Dual }, }; diff --git a/crates/openshell-server/src/auth/authenticator.rs b/crates/openshell-server/src/auth/authenticator.rs index f5d5c7b2af..a5f8c54bd2 100644 --- a/crates/openshell-server/src/auth/authenticator.rs +++ b/crates/openshell-server/src/auth/authenticator.rs @@ -14,8 +14,8 @@ //! //! Live authenticators slotting into the chain: //! - [`super::sandbox_jwt::SandboxJwtAuthenticator`] — gateway-minted JWTs -//! - [`super::k8s_sa::K8sServiceAccountAuthenticator`] — K8s projected SA -//! tokens (path-scoped to `IssueSandboxToken`) +//! - [`super::k8s_sa::SupervisorBootstrapAuthenticator`] — driver-provided +//! bootstrap tokens (path-scoped to supervisor bootstrap RPCs) //! - [`super::oidc::OidcAuthenticator`] — user OIDC Bearer tokens use super::principal::Principal; use async_trait::async_trait; diff --git a/crates/openshell-server/src/auth/guard.rs b/crates/openshell-server/src/auth/guard.rs index edcd6bc013..c1e87b6e48 100644 --- a/crates/openshell-server/src/auth/guard.rs +++ b/crates/openshell-server/src/auth/guard.rs @@ -20,6 +20,8 @@ use tracing::info; /// scope at the router level). /// - [`Principal::Sandbox`] must reference the same canonical UUID it /// was authenticated with. +/// - [`Principal::SupervisorBootstrap`] is rejected — bootstrap registration +/// identity is not a sandbox-scoped credential. /// - [`Principal::Anonymous`] is rejected — sandbox-class methods are /// never anonymously callable. /// @@ -44,6 +46,9 @@ pub fn ensure_sandbox_scope(principal: &Principal, claimed_sandbox_id: &str) -> )) } } + Principal::SupervisorBootstrap(_) => Err(Status::permission_denied( + "sandbox-scoped methods require a sandbox principal", + )), Principal::Anonymous => Err(Status::unauthenticated( "sandbox-scoped methods require an authenticated caller", )), @@ -84,7 +89,7 @@ pub fn ensure_sandbox_principal_scope( ensure_sandbox_scope(principal, claimed_sandbox_id)?; Ok(p.clone()) } - Principal::User(_) => Err(Status::permission_denied( + Principal::User(_) | Principal::SupervisorBootstrap(_) => Err(Status::permission_denied( "supervisor RPCs require a sandbox principal", )), Principal::Anonymous => Err(Status::unauthenticated( diff --git a/crates/openshell-server/src/auth/k8s_sa.rs b/crates/openshell-server/src/auth/k8s_sa.rs index eed0e5f083..fd547530ca 100644 --- a/crates/openshell-server/src/auth/k8s_sa.rs +++ b/crates/openshell-server/src/auth/k8s_sa.rs @@ -1,103 +1,59 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Kubernetes `ServiceAccount` bootstrap authenticator. +//! Supervisor bootstrap authenticator adapter. //! -//! Path-scoped to `IssueSandboxToken`. Validates a projected SA token -//! presented by a sandbox pod, reads the pod's `openshell.io/sandbox-id` -//! annotation, verifies the pod is controlled by the corresponding Sandbox CR, -//! and returns a [`Principal::Sandbox`] with -//! [`SandboxIdentitySource::K8sServiceAccount`]. The `IssueSandboxToken` handler -//! then mints a gateway-signed JWT for that sandbox id; subsequent gRPC calls -//! from the supervisor use the gateway-minted JWT validated by -//! [`super::sandbox_jwt::SandboxJwtAuthenticator`]. -//! -//! This is the only authenticator that talks to the K8s apiserver. It is -//! optional — the gateway boots without it in singleplayer deployments. +//! The Kubernetes-specific `TokenReview` and pod lookup logic lives in the +//! Kubernetes driver. This module is intentionally gateway-small: it is +//! path-scoped to supervisor bootstrap RPCs, extracts a bearer token, delegates +//! validation to the active driver's bootstrap identity provider, and turns the +//! resulting registration-only identity into the principal shape expected by the +//! gateway auth router. use super::authenticator::Authenticator; use super::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; -use async_trait::async_trait; -use k8s_openapi::api::{ - authentication::v1::{TokenReview, TokenReviewSpec, TokenReviewStatus, UserInfo}, - core::v1::Pod, +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapBinding, SupervisorBootstrapIdentityProvider, }; -use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta; -use kube::Error as KubeError; -use kube::api::{Api, ApiResource, PostParams}; -use kube::core::{DynamicObject, gvk::GroupVersionKind}; use std::sync::Arc; use tonic::Status; -use tracing::{debug, info, warn}; +use tonic::async_trait; +use tracing::{debug, warn}; -/// gRPC method path that this authenticator accepts. All other paths fall -/// through (return `Ok(None)`) so a gateway-minted JWT is required there. +/// Legacy gRPC method path accepted by this authenticator. All other +/// non-bootstrap paths fall through so a gateway-minted JWT or user credential +/// is required there. pub const ISSUE_SANDBOX_TOKEN_PATH: &str = "/openshell.v1.OpenShell/IssueSandboxToken"; +/// Supervisor registration path accepted by this authenticator. +pub const REGISTER_SUPERVISOR_POD_PATH: &str = "/openshell.v1.OpenShell/RegisterSupervisorPod"; -/// Pod annotation that binds a sandbox pod to its UUID. Set by the -/// Kubernetes compute driver at pod-create time. The gateway accepts this -/// annotation only after validating the pod's `TokenReview` binding, live UID, -/// and owning Sandbox CR. The K8s `Role` granted to the gateway must not -/// include `patch pods` (see plan §11.8). -pub const SANDBOX_ID_ANNOTATION: &str = "openshell.io/sandbox-id"; -const SANDBOX_API_GROUP: &str = "agents.x-k8s.io"; -const SANDBOX_API_VERSION_V1BETA1: &str = "v1beta1"; -const SANDBOX_API_VERSION_V1ALPHA1: &str = "v1alpha1"; -const SANDBOX_API_VERSION_FULL_V1BETA1: &str = "agents.x-k8s.io/v1beta1"; -const SANDBOX_API_VERSION_FULL_V1ALPHA1: &str = "agents.x-k8s.io/v1alpha1"; -const SANDBOX_KIND: &str = "Sandbox"; -const SANDBOX_ID_LABEL: &str = "openshell.ai/sandbox-id"; -const POD_NAME_EXTRA: &str = "authentication.kubernetes.io/pod-name"; -const POD_UID_EXTRA: &str = "authentication.kubernetes.io/pod-uid"; - -/// Resolved identity extracted from a validated SA token + pod lookup. -#[derive(Debug, Clone)] -pub struct ResolvedK8sIdentity { - pub sandbox_id: String, - pub pod_name: String, - pub pod_uid: String, -} - -/// Apiserver-facing operations the authenticator depends on. Split out so -/// tests can fake the apiserver without standing up a kube cluster. -#[async_trait] -pub trait K8sIdentityResolver: Send + Sync + 'static { - /// Validate `token` via `TokenReview` (`aud == openshell-gateway`), - /// extract the pod name/uid, then `GET` the pod and read - /// `openshell.io/sandbox-id`. Returns `Ok(None)` when the token is - /// well-formed but does not authenticate (e.g. wrong audience); returns - /// `Err` for transport/server errors. - async fn resolve(&self, token: &str) -> Result, Status>; -} - -/// Authenticator wrapper around a [`K8sIdentityResolver`]. -pub struct K8sServiceAccountAuthenticator { - resolver: Arc, +/// Path-scoped authenticator backed by the active compute driver's bootstrap +/// identity provider. +pub struct SupervisorBootstrapAuthenticator { + provider: Arc, } -impl std::fmt::Debug for K8sServiceAccountAuthenticator { +impl std::fmt::Debug for SupervisorBootstrapAuthenticator { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("K8sServiceAccountAuthenticator") + f.debug_struct("SupervisorBootstrapAuthenticator") .finish_non_exhaustive() } } -impl K8sServiceAccountAuthenticator { - pub fn new(resolver: Arc) -> Self { - Self { resolver } +impl SupervisorBootstrapAuthenticator { + pub fn new(provider: Arc) -> Self { + Self { provider } } } #[async_trait] -impl Authenticator for K8sServiceAccountAuthenticator { +impl Authenticator for SupervisorBootstrapAuthenticator { async fn authenticate( &self, headers: &http::HeaderMap, path: &str, ) -> Result, Status> { - // Scope: only the bootstrap RPC. Other paths fall through so the - // SandboxJwtAuthenticator (or OIDC) handles them. - if path != ISSUE_SANDBOX_TOKEN_PATH { + if path != ISSUE_SANDBOX_TOKEN_PATH && path != REGISTER_SUPERVISOR_POD_PATH { return Ok(None); } @@ -109,420 +65,52 @@ impl Authenticator for K8sServiceAccountAuthenticator { return Ok(None); }; - let Some(resolved) = self.resolver.resolve(token).await? else { - debug!("K8s SA token did not authenticate; falling through"); + let Some(identity) = self.provider.authenticate_registration(token).await? else { + debug!("supervisor bootstrap token did not authenticate; falling through"); return Ok(None); }; - if resolved.sandbox_id.is_empty() { + if path == REGISTER_SUPERVISOR_POD_PATH { + return Ok(Some(Principal::SupervisorBootstrap(identity))); + } + + let SupervisorBootstrapBinding::BoundSandbox { sandbox_id } = identity.binding else { warn!( - pod = %resolved.pod_name, - "pod missing openshell.io/sandbox-id annotation; rejecting" + driver = %identity.driver, + instance_name = %identity.instance_name, + instance_id = %identity.instance_id, + "bootstrap identity is not bound to a sandbox; rejecting legacy token issue" ); return Err(Status::permission_denied( - "pod is not bound to a sandbox identity", + "supervisor instance is not bound to a sandbox identity", )); - } + }; Ok(Some(Principal::Sandbox(SandboxPrincipal { - sandbox_id: resolved.sandbox_id, - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: resolved.pod_name, - pod_uid: resolved.pod_uid, + sandbox_id, + source: SandboxIdentitySource::SupervisorBootstrap { + driver: identity.driver, + instance_name: identity.instance_name, + instance_id: identity.instance_id, }, trust_domain: Some("openshell".to_string()), }))) } } -#[derive(Debug)] -struct TokenReviewIdentity { - pod_name: String, - pod_uid: String, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct SandboxOwnerReference { - api_version: String, - name: String, - uid: String, -} - -/// Resolver backed by the apiserver's `TokenReview` API and `kube::Client` -/// for the per-pod annotation lookup. -pub struct LiveK8sResolver { - token_reviews_api: Api, - pods_api: Api, - sandboxes_api_v1beta1: Api, - sandboxes_api_v1alpha1: Api, - expected_audience: String, - sandbox_namespace: String, - expected_service_account: String, -} - -impl LiveK8sResolver { - pub fn new( - client: kube::Client, - namespace: &str, - expected_audience: String, - expected_service_account: String, - ) -> Self { - let token_reviews_api: Api = Api::all(client.clone()); - let pods_api: Api = Api::namespaced(client.clone(), namespace); - let sandbox_gvk_v1beta1 = - GroupVersionKind::gvk(SANDBOX_API_GROUP, SANDBOX_API_VERSION_V1BETA1, SANDBOX_KIND); - let sandbox_resource_v1beta1 = ApiResource::from_gvk(&sandbox_gvk_v1beta1); - let sandbox_gvk_v1alpha1 = GroupVersionKind::gvk( - SANDBOX_API_GROUP, - SANDBOX_API_VERSION_V1ALPHA1, - SANDBOX_KIND, - ); - let sandbox_resource_v1alpha1 = ApiResource::from_gvk(&sandbox_gvk_v1alpha1); - let sandboxes_api_v1beta1: Api = - Api::namespaced_with(client.clone(), namespace, &sandbox_resource_v1beta1); - let sandboxes_api_v1alpha1: Api = - Api::namespaced_with(client, namespace, &sandbox_resource_v1alpha1); - Self { - token_reviews_api, - pods_api, - sandboxes_api_v1beta1, - sandboxes_api_v1alpha1, - expected_audience, - sandbox_namespace: namespace.to_string(), - expected_service_account, - } - } - - async fn get_sandbox_cr_for_owner( - &self, - owner: &SandboxOwnerReference, - ) -> Result, KubeError> { - let apis = if owner.api_version == SANDBOX_API_VERSION_FULL_V1ALPHA1 { - [&self.sandboxes_api_v1alpha1, &self.sandboxes_api_v1beta1] - } else { - [&self.sandboxes_api_v1beta1, &self.sandboxes_api_v1alpha1] - }; - - for api in apis { - match api.get_opt(&owner.name).await { - Ok(Some(sandbox_cr)) => return Ok(Some(sandbox_cr)), - Ok(None) => {} - Err(err) if should_try_next_sandbox_api_version(&err) => {} - Err(err) => return Err(err), - } - } - - Ok(None) - } -} - -#[async_trait] -impl K8sIdentityResolver for LiveK8sResolver { - async fn resolve(&self, token: &str) -> Result, Status> { - let review = TokenReview { - metadata: ObjectMeta::default(), - spec: TokenReviewSpec { - audiences: Some(vec![self.expected_audience.clone()]), - token: Some(token.to_string()), - }, - status: None, - }; - - let review = self - .token_reviews_api - .create(&PostParams::default(), &review) - .await - .map_err(|e| { - warn!(error = %e, "K8s TokenReview failed"); - Status::internal(format!("tokenreview failed: {e}")) - })?; - let status = review - .status - .ok_or_else(|| Status::internal("TokenReview response missing status"))?; - let Some(identity) = token_review_identity( - &status, - &self.expected_audience, - &self.sandbox_namespace, - &self.expected_service_account, - )? - else { - return Ok(None); - }; - - info!( - pod_name = %identity.pod_name, - pod_uid = %identity.pod_uid, - service_account = %self.expected_service_account, - "validated K8s SA token via TokenReview" - ); - - // Look up the pod and read its sandbox-id annotation. - let pod = self - .pods_api - .get_opt(&identity.pod_name) - .await - .map_err(|e| { - warn!( - pod = %identity.pod_name, - error = %e, - "failed to fetch sandbox pod for annotation lookup" - ); - Status::internal(format!("pod GET failed: {e}")) - })?; - let Some(pod) = pod else { - warn!( - pod = %identity.pod_name, - "sandbox pod referenced by SA token not found in this namespace" - ); - return Err(Status::not_found("sandbox pod not found")); - }; - - // Defense-in-depth: confirm the pod UID matches the SA token's - // `kubernetes.io.pod.uid`. Prevents a replayed token from a - // recreated pod with the same name. - let actual_uid = pod.metadata.uid.as_deref().unwrap_or_default(); - if actual_uid != identity.pod_uid { - warn!( - pod = %identity.pod_name, - claimed_uid = %identity.pod_uid, - actual_uid = %actual_uid, - "SA token pod UID does not match live pod; rejecting" - ); - return Err(Status::permission_denied("SA token pod UID mismatch")); - } - - let sandbox_id = pod_sandbox_id(&pod)?; - - let owner = sandbox_owner_reference(&pod)?; - let sandbox_cr = self.get_sandbox_cr_for_owner(&owner).await.map_err(|e| { - warn!( - pod = %identity.pod_name, - sandbox_owner = %owner.name, - sandbox_owner_api_version = %owner.api_version, - error = %e, - "failed to fetch owning Sandbox CR for pod identity validation" - ); - Status::internal(format!("sandbox GET failed: {e}")) - })?; - let Some(sandbox_cr) = sandbox_cr else { - warn!( - pod = %identity.pod_name, - sandbox_owner = %owner.name, - sandbox_owner_api_version = %owner.api_version, - "pod ownerReference points to a Sandbox CR that does not exist" - ); - return Err(Status::permission_denied("sandbox owner not found")); - }; - validate_sandbox_owner_reference(&owner, &sandbox_id, &sandbox_cr)?; - - Ok(Some(ResolvedK8sIdentity { - sandbox_id, - pod_name: identity.pod_name, - pod_uid: identity.pod_uid, - })) - } -} - -#[allow(clippy::result_large_err)] -fn token_review_identity( - status: &TokenReviewStatus, - expected_audience: &str, - sandbox_namespace: &str, - expected_service_account: &str, -) -> Result, Status> { - if status.authenticated != Some(true) { - debug!( - error = status.error.as_deref().unwrap_or_default(), - "K8s TokenReview did not authenticate token" - ); - return Ok(None); - } - - let audiences = status.audiences.as_deref().unwrap_or_default(); - if !audiences.iter().any(|aud| aud == expected_audience) { - warn!( - expected_audience = %expected_audience, - audiences = ?audiences, - "K8s TokenReview authenticated token without expected audience" - ); - return Err(Status::unauthenticated("SA token audience not accepted")); - } - - let user = status - .user - .as_ref() - .ok_or_else(|| Status::permission_denied("TokenReview response missing user info"))?; - let username = user - .username - .as_deref() - .ok_or_else(|| Status::permission_denied("TokenReview response missing username"))?; - let expected_username = - format!("system:serviceaccount:{sandbox_namespace}:{expected_service_account}"); - if username != expected_username { - warn!( - username = %username, - sandbox_namespace = %sandbox_namespace, - service_account = %expected_service_account, - "K8s TokenReview principal is not the configured sandbox service account" - ); - return Err(Status::permission_denied( - "SA token is not from the configured sandbox service account", - )); - } - - let pod_name = user_extra_one(user, POD_NAME_EXTRA)?; - let pod_uid = user_extra_one(user, POD_UID_EXTRA)?; - Ok(Some(TokenReviewIdentity { pod_name, pod_uid })) -} - -#[allow(clippy::result_large_err)] -fn user_extra_one(user: &UserInfo, key: &str) -> Result { - let Some(values) = user.extra.as_ref().and_then(|extra| extra.get(key)) else { - return Err(Status::permission_denied("SA token is not pod-bound")); - }; - if values.len() != 1 || values[0].is_empty() { - return Err(Status::permission_denied( - "SA token has invalid pod binding", - )); - } - Ok(values[0].clone()) -} - -#[allow(clippy::result_large_err)] -fn pod_sandbox_id(pod: &Pod) -> Result { - let sandbox_id = pod - .metadata - .annotations - .as_ref() - .and_then(|a| a.get(SANDBOX_ID_ANNOTATION)) - .cloned() - .unwrap_or_default(); - if sandbox_id.is_empty() { - return Err(Status::permission_denied( - "pod is not bound to a sandbox identity", - )); - } - Ok(sandbox_id) -} - -#[allow(clippy::result_large_err)] -fn sandbox_owner_reference(pod: &Pod) -> Result { - let owner_refs = pod.metadata.owner_references.as_deref().unwrap_or_default(); - let mut sandbox_refs = owner_refs - .iter() - .filter(|owner| is_supported_sandbox_owner_reference(owner)); - let Some(owner) = sandbox_refs.next() else { - let unsupported_sandbox_api_versions = owner_refs - .iter() - .filter(|owner| owner.kind == SANDBOX_KIND) - .map(|owner| owner.api_version.as_str()) - .collect::>(); - if !unsupported_sandbox_api_versions.is_empty() { - warn!( - api_versions = ?unsupported_sandbox_api_versions, - supported_api_versions = ?[ - SANDBOX_API_VERSION_FULL_V1BETA1, - SANDBOX_API_VERSION_FULL_V1ALPHA1, - ], - "pod Sandbox ownerReference uses unsupported apiVersion" - ); - } - return Err(Status::permission_denied( - "pod is not controlled by an OpenShell Sandbox", - )); - }; - if sandbox_refs.next().is_some() { - return Err(Status::permission_denied( - "pod has multiple OpenShell Sandbox owners", - )); - } - if owner.controller != Some(true) { - return Err(Status::permission_denied( - "pod Sandbox ownerReference is not controlling", - )); - } - if owner.name.is_empty() || owner.uid.is_empty() { - return Err(Status::permission_denied( - "pod Sandbox ownerReference is incomplete", - )); - } - Ok(SandboxOwnerReference { - api_version: owner.api_version.clone(), - name: owner.name.clone(), - uid: owner.uid.clone(), - }) -} - -fn is_supported_sandbox_owner_reference( - owner: &k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference, -) -> bool { - owner.kind == SANDBOX_KIND - && matches!( - owner.api_version.as_str(), - SANDBOX_API_VERSION_FULL_V1BETA1 | SANDBOX_API_VERSION_FULL_V1ALPHA1 - ) -} - -fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { - // Kubernetes returns a structured 404 for some missing API resources and a - // raw "404 page not found" body for others. Both mean the probed - // group/version is unavailable and the next supported Sandbox API version - // should be tried. - matches!(err, KubeError::Api(api) if api.code == 404) -} - -#[allow(clippy::result_large_err)] -fn validate_sandbox_owner_reference( - owner: &SandboxOwnerReference, - sandbox_id: &str, - sandbox_cr: &DynamicObject, -) -> Result<(), Status> { - let actual_uid = sandbox_cr.metadata.uid.as_deref().unwrap_or_default(); - if actual_uid != owner.uid { - warn!( - sandbox_owner = %owner.name, - owner_uid = %owner.uid, - actual_uid = %actual_uid, - "pod Sandbox ownerReference UID does not match live Sandbox CR" - ); - return Err(Status::permission_denied("sandbox owner UID mismatch")); - } - - let actual_sandbox_id = sandbox_cr - .metadata - .labels - .as_ref() - .and_then(|labels| labels.get(SANDBOX_ID_LABEL)) - .map(String::as_str) - .unwrap_or_default(); - if actual_sandbox_id != sandbox_id { - warn!( - sandbox_owner = %owner.name, - owner_uid = %owner.uid, - pod_sandbox_id = %sandbox_id, - cr_sandbox_id = %actual_sandbox_id, - "pod sandbox annotation does not match owning Sandbox CR label" - ); - return Err(Status::permission_denied("sandbox owner ID mismatch")); - } - - Ok(()) -} - #[cfg(test)] -pub mod test_support { +mod tests { use super::*; + use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentity; use std::sync::Mutex; - /// Fake resolver for unit tests. Returns the configured outcome on - /// every call and records the tokens it observed. - pub struct FakeResolver { - pub outcome: Result, Status>, - pub seen_tokens: Mutex>, + struct FakeProvider { + outcome: Result, Status>, + seen_tokens: Mutex>, } - impl FakeResolver { - pub fn returning(outcome: Result, Status>) -> Self { + impl FakeProvider { + fn returning(outcome: Result, Status>) -> Self { Self { outcome, seen_tokens: Mutex::new(Vec::new()), @@ -531,403 +119,89 @@ pub mod test_support { } #[async_trait] - impl K8sIdentityResolver for FakeResolver { - async fn resolve(&self, token: &str) -> Result, Status> { + impl SupervisorBootstrapIdentityProvider for FakeProvider { + async fn authenticate_registration( + &self, + token: &str, + ) -> Result, Status> { self.seen_tokens.lock().unwrap().push(token.to_string()); match &self.outcome { - Ok(opt) => Ok(opt.clone()), - Err(s) => Err(Status::new(s.code(), s.message())), + Ok(identity) => Ok(identity.clone()), + Err(status) => Err(Status::new(status.code(), status.message())), } } } -} - -#[cfg(test)] -mod tests { - use super::test_support::FakeResolver; - use super::*; - use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference; - use std::collections::BTreeMap; fn bearer_headers(token: &str) -> http::HeaderMap { - let mut h = http::HeaderMap::new(); - h.insert( + let mut headers = http::HeaderMap::new(); + headers.insert( "authorization", http::HeaderValue::from_str(&format!("Bearer {token}")).unwrap(), ); - h - } - - fn kube_api_error(code: u16, message: &str) -> KubeError { - KubeError::Api(kube::core::ErrorResponse { - status: if code == 404 { - "404 Not Found".to_string() - } else { - "Failure".to_string() - }, - message: message.to_string(), - reason: "Failed to parse error data".to_string(), - code, - }) + headers } - #[test] - fn sandbox_api_version_probe_retries_on_structured_and_raw_404() { - let structured = kube_api_error(404, "could not find the requested resource"); - assert!(should_try_next_sandbox_api_version(&structured)); - - let raw = kube_api_error(404, "404 page not found\n"); - assert!(should_try_next_sandbox_api_version(&raw)); - } - - #[test] - fn sandbox_api_version_probe_keeps_non_404_errors() { - let err = kube_api_error(403, "sandboxes.agents.x-k8s.io is forbidden"); - assert!(!should_try_next_sandbox_api_version(&err)); - } - - fn token_review_status( - authenticated: bool, - audiences: Vec<&str>, - username: &str, - extra: Vec<(&str, &str)>, - ) -> TokenReviewStatus { - TokenReviewStatus { - authenticated: Some(authenticated), - audiences: Some(audiences.into_iter().map(str::to_string).collect()), - error: None, - user: Some(UserInfo { - username: Some(username.to_string()), - uid: Some("sa-uid".to_string()), - groups: Some(vec![ - "system:serviceaccounts".to_string(), - "system:serviceaccounts:openshell".to_string(), - "system:authenticated".to_string(), - ]), - extra: Some( - extra - .into_iter() - .map(|(k, v)| (k.to_string(), vec![v.to_string()])) - .collect::>(), - ), - }), + fn identity(binding: SupervisorBootstrapBinding) -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "openshell-sandbox-a".to_string(), + instance_id: "uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "cr-uid-a".to_string(), + binding, } } - fn sandbox_owner(name: &str, uid: &str) -> OwnerReference { - sandbox_owner_with_api_version(SANDBOX_API_VERSION_FULL_V1BETA1, name, uid) - } - - fn sandbox_owner_with_api_version(api_version: &str, name: &str, uid: &str) -> OwnerReference { - OwnerReference { - api_version: api_version.to_string(), - block_owner_deletion: None, - controller: Some(true), - kind: SANDBOX_KIND.to_string(), - name: name.to_string(), - uid: uid.to_string(), - } - } - - fn pod_with_owner_refs(owner_references: Vec) -> Pod { - Pod { - metadata: ObjectMeta { - owner_references: Some(owner_references), - ..Default::default() - }, - ..Default::default() - } - } - - fn pod_with_sandbox_id(sandbox_id: Option<&str>) -> Pod { - Pod { - metadata: ObjectMeta { - annotations: sandbox_id.map(|id| { - BTreeMap::from([(SANDBOX_ID_ANNOTATION.to_string(), id.to_string())]) - }), - ..Default::default() - }, - ..Default::default() - } - } - - fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str) -> DynamicObject { - let sandbox_gvk = - GroupVersionKind::gvk(SANDBOX_API_GROUP, SANDBOX_API_VERSION_V1BETA1, SANDBOX_KIND); - let sandbox_resource = ApiResource::from_gvk(&sandbox_gvk); - let mut cr = DynamicObject::new(name, &sandbox_resource); - cr.metadata.uid = Some(uid.to_string()); - cr.metadata.labels = Some(BTreeMap::from([( - SANDBOX_ID_LABEL.to_string(), - sandbox_id.to_string(), - )])); - cr - } - - #[test] - fn token_review_identity_extracts_pod_binding() { - let status = token_review_status( - true, - vec!["openshell-gateway"], - "system:serviceaccount:openshell:default", - vec![ - (POD_NAME_EXTRA, "openshell-sandbox-a"), - (POD_UID_EXTRA, "uid-a"), - ], - ); - - let identity = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .unwrap() - .expect("authenticated token should resolve"); - - assert_eq!(identity.pod_name, "openshell-sandbox-a"); - assert_eq!(identity.pod_uid, "uid-a"); - } - - #[test] - fn token_review_identity_returns_none_when_not_authenticated() { - let status = TokenReviewStatus { - authenticated: Some(false), - error: Some("invalid audience".to_string()), - ..Default::default() - }; - - assert!( - token_review_identity(&status, "openshell-gateway", "openshell", "default") - .unwrap() - .is_none() - ); - } - - #[test] - fn token_review_identity_requires_expected_audience() { - let status = token_review_status( - true, - vec!["kubernetes.default.svc"], - "system:serviceaccount:openshell:default", - vec![ - (POD_NAME_EXTRA, "openshell-sandbox-a"), - (POD_UID_EXTRA, "uid-a"), - ], - ); - - let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .expect_err("wrong audience must fail closed"); - assert_eq!(err.code(), tonic::Code::Unauthenticated); - } - - #[test] - fn token_review_identity_requires_sandbox_namespace() { - let status = token_review_status( - true, - vec!["openshell-gateway"], - "system:serviceaccount:other:default", - vec![ - (POD_NAME_EXTRA, "openshell-sandbox-a"), - (POD_UID_EXTRA, "uid-a"), - ], - ); - - let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .expect_err("other namespace must be rejected"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn token_review_identity_requires_configured_service_account() { - let status = token_review_status( - true, - vec!["openshell-gateway"], - "system:serviceaccount:openshell:other", - vec![ - (POD_NAME_EXTRA, "openshell-sandbox-a"), - (POD_UID_EXTRA, "uid-a"), - ], - ); - - let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .expect_err("other service account must be rejected"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn token_review_identity_requires_pod_bound_extras() { - let status = token_review_status( - true, - vec!["openshell-gateway"], - "system:serviceaccount:openshell:default", - vec![], - ); - - let err = token_review_identity(&status, "openshell-gateway", "openshell", "default") - .expect_err("non pod-bound tokens must be rejected"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn pod_sandbox_id_requires_annotation() { - assert_eq!( - pod_sandbox_id(&pod_with_sandbox_id(Some("sandbox-id-a"))).unwrap(), - "sandbox-id-a" - ); - - let err = pod_sandbox_id(&pod_with_sandbox_id(None)) - .expect_err("missing sandbox-id annotation must fail"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn sandbox_owner_reference_extracts_controlling_sandbox_owner() { - let pod = pod_with_owner_refs(vec![sandbox_owner("sandbox-a", "cr-uid-a")]); - - let owner = sandbox_owner_reference(&pod).expect("expected Sandbox owner"); - - assert_eq!( - owner, - SandboxOwnerReference { - api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), - name: "sandbox-a".to_string(), - uid: "cr-uid-a".to_string(), - } - ); - } - - #[test] - fn sandbox_owner_reference_accepts_v1alpha1_owner() { - let pod = pod_with_owner_refs(vec![sandbox_owner_with_api_version( - SANDBOX_API_VERSION_FULL_V1ALPHA1, - "sandbox-a", - "cr-uid-a", - )]); - - let owner = sandbox_owner_reference(&pod).expect("expected v1alpha1 Sandbox owner"); - - assert_eq!( - owner, - SandboxOwnerReference { - api_version: SANDBOX_API_VERSION_FULL_V1ALPHA1.to_string(), - name: "sandbox-a".to_string(), - uid: "cr-uid-a".to_string(), - } - ); - } - - #[test] - fn sandbox_owner_reference_rejects_missing_owner() { - let pod = pod_with_owner_refs(vec![]); - - let err = sandbox_owner_reference(&pod).expect_err("missing owner must fail"); - - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn sandbox_owner_reference_rejects_unsupported_sandbox_api_version() { - let pod = pod_with_owner_refs(vec![sandbox_owner_with_api_version( - "agents.x-k8s.io/v1", - "sandbox-a", - "cr-uid-a", - )]); - - let err = - sandbox_owner_reference(&pod).expect_err("unsupported apiVersion must fail closed"); - - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn sandbox_owner_reference_requires_controlling_owner() { - let mut owner = sandbox_owner("sandbox-a", "cr-uid-a"); - owner.controller = Some(false); - let pod = pod_with_owner_refs(vec![owner]); - - let err = sandbox_owner_reference(&pod).expect_err("non-controller owner must fail"); - - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn sandbox_owner_reference_rejects_ambiguous_sandbox_owners() { - let pod = pod_with_owner_refs(vec![ - sandbox_owner("sandbox-a", "cr-uid-a"), - sandbox_owner("sandbox-b", "cr-uid-b"), - ]); - - let err = sandbox_owner_reference(&pod).expect_err("multiple owners must fail"); - - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - - #[test] - fn validate_sandbox_owner_reference_requires_matching_cr_uid_and_label() { - let owner = SandboxOwnerReference { - api_version: SANDBOX_API_VERSION_FULL_V1BETA1.to_string(), - name: "sandbox-a".to_string(), - uid: "cr-uid-a".to_string(), - }; - let cr = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-a"); - validate_sandbox_owner_reference(&owner, "sandbox-id-a", &cr) - .expect("matching CR should be accepted"); - - let wrong_uid = sandbox_cr("sandbox-a", "cr-uid-b", "sandbox-id-a"); - let err = validate_sandbox_owner_reference(&owner, "sandbox-id-a", &wrong_uid) - .expect_err("wrong CR UID must fail"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - - let wrong_label = sandbox_cr("sandbox-a", "cr-uid-a", "sandbox-id-b"); - let err = validate_sandbox_owner_reference(&owner, "sandbox-id-a", &wrong_label) - .expect_err("wrong sandbox-id label must fail"); - assert_eq!(err.code(), tonic::Code::PermissionDenied); - } - #[tokio::test] - async fn authenticates_on_issue_path_only() { - let resolved = ResolvedK8sIdentity { - sandbox_id: "sandbox-a".to_string(), - pod_name: "openshell-sandbox-a".to_string(), - pod_uid: "uid-a".to_string(), - }; - let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); - let auth = K8sServiceAccountAuthenticator::new(fake.clone()); + async fn authenticates_on_bootstrap_paths_only() { + let provider = Arc::new(FakeProvider::returning(Ok(Some(identity( + SupervisorBootstrapBinding::BoundSandbox { + sandbox_id: "sandbox-a".to_string(), + }, + ))))); + let auth = SupervisorBootstrapAuthenticator::new(provider.clone()); - let on_issue = auth - .authenticate(&bearer_headers("sa-jwt"), ISSUE_SANDBOX_TOKEN_PATH) + let issue = auth + .authenticate(&bearer_headers("bootstrap-token"), ISSUE_SANDBOX_TOKEN_PATH) .await .unwrap() .expect("expected principal"); - match on_issue { + match issue { Principal::Sandbox(p) => { assert_eq!(p.sandbox_id, "sandbox-a"); assert!(matches!( p.source, - SandboxIdentitySource::K8sServiceAccount { .. } + SandboxIdentitySource::SupervisorBootstrap { .. } )); } _ => panic!("expected sandbox principal"), } - let off_issue = auth + let register = auth + .authenticate( + &bearer_headers("bootstrap-token"), + REGISTER_SUPERVISOR_POD_PATH, + ) + .await + .unwrap() + .expect("expected principal"); + assert!(matches!(register, Principal::SupervisorBootstrap(_))); + + let off_path = auth .authenticate( - &bearer_headers("sa-jwt"), + &bearer_headers("bootstrap-token"), "/openshell.v1.OpenShell/GetSandboxConfig", ) .await .unwrap(); - assert!( - off_issue.is_none(), - "K8s SA authenticator must be scoped to IssueSandboxToken" - ); - assert_eq!( - fake.seen_tokens.lock().unwrap().len(), - 1, - "off-path call must not consult the apiserver" - ); + assert!(off_path.is_none()); + assert_eq!(provider.seen_tokens.lock().unwrap().len(), 2); } #[tokio::test] async fn missing_bearer_yields_none() { - let fake = Arc::new(FakeResolver::returning(Ok(None))); - let auth = K8sServiceAccountAuthenticator::new(fake); + let provider = Arc::new(FakeProvider::returning(Ok(None))); + let auth = SupervisorBootstrapAuthenticator::new(provider); let result = auth .authenticate(&http::HeaderMap::new(), ISSUE_SANDBOX_TOKEN_PATH) .await @@ -936,45 +210,39 @@ mod tests { } #[tokio::test] - async fn resolver_returning_none_falls_through() { - let fake = Arc::new(FakeResolver::returning(Ok(None))); - let auth = K8sServiceAccountAuthenticator::new(fake); - let result = auth - .authenticate( - &bearer_headers("not-a-real-sa-token"), - ISSUE_SANDBOX_TOKEN_PATH, - ) - .await - .unwrap(); - assert!(result.is_none(), "non-authenticating tokens fall through"); - } - - #[tokio::test] - async fn pod_without_annotation_is_rejected() { - let resolved = ResolvedK8sIdentity { - sandbox_id: String::new(), - pod_name: "stray-pod".to_string(), - pod_uid: "uid".to_string(), - }; - let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); - let auth = K8sServiceAccountAuthenticator::new(fake); + async fn issue_token_rejects_unbound_identity() { + let provider = Arc::new(FakeProvider::returning(Ok(Some(identity( + SupervisorBootstrapBinding::WarmPending, + ))))); + let auth = SupervisorBootstrapAuthenticator::new(provider); let err = auth - .authenticate(&bearer_headers("sa-jwt"), ISSUE_SANDBOX_TOKEN_PATH) + .authenticate(&bearer_headers("bootstrap-token"), ISSUE_SANDBOX_TOKEN_PATH) .await - .expect_err("unbound pod must be rejected"); + .expect_err("unbound identity must be rejected"); assert_eq!(err.code(), tonic::Code::PermissionDenied); } #[tokio::test] - async fn resolver_error_propagates() { - let fake = Arc::new(FakeResolver::returning(Err(Status::unavailable( - "apiserver down", - )))); - let auth = K8sServiceAccountAuthenticator::new(fake); - let err = auth - .authenticate(&bearer_headers("sa-jwt"), ISSUE_SANDBOX_TOKEN_PATH) + async fn register_accepts_unbound_identity() { + let provider = Arc::new(FakeProvider::returning(Ok(Some(identity( + SupervisorBootstrapBinding::WarmPending, + ))))); + let auth = SupervisorBootstrapAuthenticator::new(provider); + let principal = auth + .authenticate( + &bearer_headers("bootstrap-token"), + REGISTER_SUPERVISOR_POD_PATH, + ) .await - .expect_err("resolver error must propagate"); - assert_eq!(err.code(), tonic::Code::Unavailable); + .unwrap() + .expect("expected principal"); + + match principal { + Principal::SupervisorBootstrap(identity) => { + assert_eq!(identity.binding, SupervisorBootstrapBinding::WarmPending); + assert_eq!(identity.owner_uid, "cr-uid-a"); + } + _ => panic!("expected bootstrap principal"), + } } } diff --git a/crates/openshell-server/src/auth/method_authz.rs b/crates/openshell-server/src/auth/method_authz.rs index ec8dc5bca3..c53ca8024d 100644 --- a/crates/openshell-server/src/auth/method_authz.rs +++ b/crates/openshell-server/src/auth/method_authz.rs @@ -34,6 +34,9 @@ pub enum AuthMode { /// Only callable by a `Principal::Sandbox` (gateway-minted sandbox JWT). /// See `auth/sandbox_jwt.rs`. Sandbox, + /// Only callable by a Kubernetes pod principal during supervisor pod + /// registration. This does not grant sandbox-scoped RPC access. + PodRegistration, /// Bearer (OIDC) authentication required. Bearer, /// Either sandbox principal or Bearer; scope and role apply on @@ -109,6 +112,16 @@ pub fn is_sandbox_callable(method: &str) -> bool { ) } +/// `true` if the method is callable by a Kubernetes pod registration +/// principal. +#[must_use] +pub fn is_pod_registration_callable(method: &str) -> bool { + matches!( + lookup(method).map(|m| m.mode), + Some(AuthMode::PodRegistration) + ) +} + /// `true` if the method is callable by a `Principal::User` (`bearer` or /// `dual` auth mode). /// @@ -119,7 +132,7 @@ pub fn is_sandbox_callable(method: &str) -> bool { #[must_use] pub fn is_user_callable(method: &str) -> bool { match lookup(method).map(|m| m.mode) { - Some(AuthMode::Sandbox | AuthMode::Unauthenticated) => false, + Some(AuthMode::Sandbox | AuthMode::PodRegistration | AuthMode::Unauthenticated) => false, Some(AuthMode::Bearer | AuthMode::Dual) | None => true, } } diff --git a/crates/openshell-server/src/auth/principal.rs b/crates/openshell-server/src/auth/principal.rs index 1d4cb7276c..536bb08703 100644 --- a/crates/openshell-server/src/auth/principal.rs +++ b/crates/openshell-server/src/auth/principal.rs @@ -15,6 +15,7 @@ //! to prevent cross-sandbox access (see issue #1354). use super::identity::Identity; +use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentity; /// Who is calling. /// @@ -28,6 +29,13 @@ pub enum Principal { /// sandbox UUID. The wrapped `sandbox_id` MUST match any sandbox referenced /// in the request body for sandbox-class methods. Sandbox(#[allow(dead_code)] SandboxPrincipal), + /// Driver runtime instance authenticated for supervisor bootstrap + /// registration only. + /// + /// This is intentionally not a sandbox principal: warm pods can be valid + /// driver instances before they are claimed by a sandbox. The router only + /// allows this principal to call `RegisterSupervisorPod`. + SupervisorBootstrap(SupervisorBootstrapIdentity), /// Truly unauthenticated caller (health probes, reflection). Sandbox-class /// and user-class methods reject this variant. #[allow(dead_code)] @@ -70,7 +78,11 @@ pub enum SandboxIdentitySource { /// Per-sandbox client certificate. Reserved for channel-bound sandbox /// identity. BootstrapCert { fingerprint: String }, - /// K8s `ServiceAccount` token used to bootstrap a gateway-minted JWT - /// via `IssueSandboxToken`. Populated only on that one RPC path. - K8sServiceAccount { pod_name: String, pod_uid: String }, + /// Driver-provided bootstrap token used to bootstrap a gateway-minted JWT. + /// Populated only on supervisor bootstrap RPC paths. + SupervisorBootstrap { + driver: String, + instance_name: String, + instance_id: String, + }, } diff --git a/crates/openshell-server/src/auth/sandbox_jwt.rs b/crates/openshell-server/src/auth/sandbox_jwt.rs index 39f5982ca0..46dab8e112 100644 --- a/crates/openshell-server/src/auth/sandbox_jwt.rs +++ b/crates/openshell-server/src/auth/sandbox_jwt.rs @@ -8,7 +8,7 @@ //! supervisor-to-gateway gRPC calls. This module implements both sides of the //! gateway-controlled token: //! - [`SandboxJwtIssuer`] mints fresh tokens (called from -//! `handle_create_sandbox` and the `IssueSandboxToken` RPC). +//! `handle_create_sandbox` and the Kubernetes bootstrap RPCs). //! - [`SandboxJwtAuthenticator`] validates tokens on inbound requests and //! produces a [`Principal::Sandbox`] with [`SandboxIdentitySource::BootstrapJwt`]. //! diff --git a/crates/openshell-server/src/auth/sandbox_methods.rs b/crates/openshell-server/src/auth/sandbox_methods.rs index b90841d85a..2566c29ab5 100644 --- a/crates/openshell-server/src/auth/sandbox_methods.rs +++ b/crates/openshell-server/src/auth/sandbox_methods.rs @@ -51,6 +51,9 @@ mod tests { assert!(!is_sandbox_callable( "/openshell.inference.v1.Inference/GetInferenceRoute" )); + assert!(!is_sandbox_callable( + "/openshell.v1.OpenShell/RegisterSupervisorPod" + )); assert!(!is_sandbox_callable( "/openshell.inference.v1.Inference/SetInferenceRoute" )); diff --git a/crates/openshell-server/src/compute/driver_config.rs b/crates/openshell-server/src/compute/driver_config.rs index 59fed439bc..1d13277074 100644 --- a/crates/openshell-server/src/compute/driver_config.rs +++ b/crates/openshell-server/src/compute/driver_config.rs @@ -54,22 +54,6 @@ pub fn kubernetes_config_from_context( Ok(cfg) } -pub fn kubernetes_config_for_k8s_sa_bootstrap( - file: Option<&config_file::ConfigFile>, -) -> Result { - let Some(file) = file else { - return Err(Error::config( - "K8s ServiceAccount bootstrap requires [openshell.drivers.kubernetes] when sandbox JWT issuing is enabled in-cluster", - )); - }; - if !file.openshell.drivers.contains_key("kubernetes") { - return Err(Error::config( - "K8s ServiceAccount bootstrap requires [openshell.drivers.kubernetes] when sandbox JWT issuing is enabled in-cluster", - )); - } - driver_config_from_file(Some(file), ComputeDriverKind::Kubernetes.as_str()) -} - /// Build the selected Podman config from TOML plus runtime defaults. pub fn podman_config_from_context( context: DriverStartupContext<'_>, @@ -262,35 +246,6 @@ mod tests { } } - #[test] - fn k8s_sa_bootstrap_rejects_missing_kubernetes_driver_config() { - let err = kubernetes_config_for_k8s_sa_bootstrap(None).unwrap_err(); - assert!(err.to_string().contains("[openshell.drivers.kubernetes]")); - - let file: config_file::ConfigFile = - toml::from_str("[openshell.gateway]\n").expect("valid config"); - let err = kubernetes_config_for_k8s_sa_bootstrap(Some(&file)).unwrap_err(); - assert!(err.to_string().contains("[openshell.drivers.kubernetes]")); - } - - #[test] - fn k8s_sa_bootstrap_uses_configured_namespace_and_service_account() { - let file: config_file::ConfigFile = toml::from_str( - r#" -[openshell.gateway] - -[openshell.drivers.kubernetes] -namespace = "sandboxes" -service_account_name = "sandbox-sa" -"#, - ) - .expect("valid config"); - - let cfg = kubernetes_config_for_k8s_sa_bootstrap(Some(&file)).unwrap(); - assert_eq!(cfg.namespace, "sandboxes"); - assert_eq!(cfg.service_account_name, "sandbox-sa"); - } - #[test] fn podman_config_reads_bind_mount_opt_in_from_driver_table() { let file: config_file::ConfigFile = toml::from_str( diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index cf6b17c7e9..b945d44fb7 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -37,10 +37,15 @@ use openshell_core::proto::{ PlatformEvent, Sandbox, SandboxCondition, SandboxPhase, SandboxSpec, SandboxStatus, SandboxTemplate, ServiceEndpoint, SshSession, }; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapActivator, SupervisorBootstrapIdentityProvider, +}; use openshell_core::{ObjectLabels, ObjectWorkspace}; use openshell_driver_docker::DockerComputeDriver; use openshell_driver_kubernetes::{ ComputeDriverService as KubernetesDriverService, KubernetesComputeDriver, + KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, + SandboxClaimActivationController, }; use openshell_driver_podman::{ComputeDriverService as PodmanDriverService, PodmanComputeDriver}; use prost::Message; @@ -231,6 +236,54 @@ impl Drop for ManagedDriverProcess { } } +async fn kubernetes_supervisor_bootstrap_identity_provider( + config: &KubernetesComputeConfig, +) -> Option> { + std::env::var_os("KUBERNETES_SERVICE_HOST")?; + + match kube::Client::try_default().await { + Ok(client) => { + let resolver = Arc::new(LiveK8sResolver::new( + client, + &config.namespace, + "openshell-gateway".to_string(), + config.service_account_name.clone(), + )); + Some(Arc::new( + KubernetesSupervisorBootstrapIdentityProvider::new(resolver), + )) + } + Err(err) => { + warn!( + error = %err, + "in-cluster K8s client construction failed; supervisor bootstrap is disabled" + ); + None + } + } +} + +async fn kubernetes_sandbox_claim_activation_controller( + config: &KubernetesComputeConfig, +) -> Option> { + std::env::var_os("KUBERNETES_SERVICE_HOST")?; + + match kube::Client::try_default().await { + Ok(client) => Some(Arc::new(SandboxClaimActivationController::new( + client.clone(), + client, + config.namespace.clone(), + ))), + Err(err) => { + warn!( + error = %err, + "in-cluster K8s client construction failed; SandboxClaim activation is disabled" + ); + None + } + } +} + #[derive(Debug)] pub struct AcquiredRemoteDriverEndpoint { pub(crate) name: String, @@ -368,6 +421,8 @@ pub struct ComputeRuntime { sandbox_watch_bus: SandboxWatchBus, tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, + supervisor_bootstrap_identity: Option>, + sandbox_claim_activation: Option>, sync_lock: Arc>, delete_gates: Arc, gateway_bind_addresses: Vec, @@ -393,6 +448,8 @@ impl ComputeRuntime { sandbox_watch_bus: SandboxWatchBus, tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, + supervisor_bootstrap_identity: Option>, + sandbox_claim_activation: Option>, gateway_bind_addresses: Vec, ) -> Result { let capabilities = driver @@ -425,6 +482,8 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + supervisor_bootstrap_identity, + sandbox_claim_activation, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses, @@ -486,6 +545,8 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + None, + None, gateway_bind_addresses, ) .await @@ -499,6 +560,10 @@ impl ComputeRuntime { tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, ) -> Result { + let supervisor_bootstrap_identity = + kubernetes_supervisor_bootstrap_identity_provider(&config).await; + let sandbox_claim_activation = + kubernetes_sandbox_claim_activation_controller(&config).await; let driver = KubernetesComputeDriver::new(config) .await .map_err(|err| ComputeError::Message(err.to_string()))?; @@ -514,6 +579,8 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + supervisor_bootstrap_identity, + sandbox_claim_activation, Vec::new(), ) .await @@ -539,6 +606,8 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + None, + None, Vec::new(), ) .await @@ -567,6 +636,8 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + None, + None, Vec::new(), ) .await @@ -582,6 +653,23 @@ impl ComputeRuntime { std::slice::from_ref(&self.driver_info) } + #[must_use] + pub fn supervisor_bootstrap_identity_provider( + &self, + ) -> Option> { + self.supervisor_bootstrap_identity.clone() + } + + pub(crate) fn spawn_sandbox_claim_activation( + &self, + activator: Arc, + shutdown_rx: watch::Receiver, + ) { + if let Some(controller) = self.sandbox_claim_activation.clone() { + controller.spawn(activator, shutdown_rx); + } + } + #[must_use] pub fn driver_kind(&self) -> Option { self.driver_info.name.parse().ok() @@ -662,7 +750,43 @@ impl ComputeRuntime { })) .await { - Ok(_) => { + Ok(response) => { + let activation_token = response.into_inner().activation_token; + if !activation_token.is_empty() + && let Err(err) = self + .store + .put_activation_token(sandbox.object_id(), &activation_token) + .await + { + warn!( + sandbox_id = %sandbox.object_id(), + error = %err, + "Failed to persist warm-pool activation token mapping after driver create" + ); + match self + .driver + .delete_sandbox(Request::new(DeleteSandboxRequest { + sandbox_id: sandbox.object_id().to_string(), + sandbox_name: sandbox.object_name().to_string(), + })) + .await + { + Ok(_) => {} + Err(cleanup_err) => warn!( + sandbox_id = %sandbox.object_id(), + error = %cleanup_err.message(), + "Failed to clean up driver resource after activation token persistence failure" + ), + } + let _ = self + .store + .delete(Sandbox::object_type(), sandbox.object_id()) + .await; + self.sandbox_index.remove_sandbox(sandbox.object_id()); + return Err(Status::internal( + "persist activation token mapping failed after driver create", + )); + } self.sandbox_watch_bus.notify(sandbox.object_id()); if let Some(metadata) = sandbox.metadata.as_mut() { metadata.resource_version = result.resource_version; @@ -798,6 +922,8 @@ impl ComputeRuntime { if deleted { self.cleanup_local_state_if_sandbox_absent(&delete_guard, &target.sandbox_id) .await?; + self.cleanup_activation_token_mapping(&target.sandbox_id) + .await; } else if !self .remove_deleting_sandbox_record(&delete_guard, &target.sandbox_id) .await @@ -971,6 +1097,7 @@ impl ComputeRuntime { .map_err(|err| err.to_string())? else { self.cleanup_removed_sandbox_state(sandbox_id); + self.cleanup_activation_token_mapping(sandbox_id).await; return Ok(false); }; @@ -998,10 +1125,12 @@ impl ComputeRuntime { { Ok(true) => { self.cleanup_removed_sandbox_state(sandbox_id); + self.cleanup_activation_token_mapping(sandbox_id).await; Ok(true) } Ok(false) => { self.cleanup_removed_sandbox_state(sandbox_id); + self.cleanup_activation_token_mapping(sandbox_id).await; Ok(false) } Err(crate::persistence::PersistenceError::Conflict { @@ -1917,10 +2046,25 @@ impl ComputeRuntime { .map_err(|err| Status::internal(format!("fetch sandbox failed: {err}")))?; if record.is_none() { self.cleanup_removed_sandbox_state(sandbox_id); + self.cleanup_activation_token_mapping(sandbox_id).await; } Ok(()) } + async fn cleanup_activation_token_mapping(&self, sandbox_id: &str) { + if let Err(err) = self + .store + .delete_activation_token_for_sandbox(sandbox_id) + .await + { + warn!( + sandbox_id, + error = %err, + "Failed to delete warm-pool activation token mapping" + ); + } + } + fn cleanup_removed_sandbox_state(&self, sandbox_id: &str) { self.sandbox_index.remove_sandbox(sandbox_id); self.sandbox_watch_bus.notify(sandbox_id); @@ -2709,7 +2853,9 @@ impl ComputeDriver for NoopTestDriver { ) -> Result, Status> { Ok(tonic::Response::new( - openshell_core::proto::compute::v1::CreateSandboxResponse {}, + openshell_core::proto::compute::v1::CreateSandboxResponse { + activation_token: String::new(), + }, )) } @@ -2759,6 +2905,8 @@ pub async fn new_test_runtime(store: Arc) -> ComputeRuntime { sandbox_watch_bus: SandboxWatchBus::new(), tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), + supervisor_bootstrap_identity: None, + sandbox_claim_activation: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses: Vec::new(), @@ -2970,7 +3118,9 @@ mod tests { &self, _request: Request, ) -> Result, Status> { - Ok(tonic::Response::new(CreateSandboxResponse {})) + Ok(tonic::Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( @@ -3150,7 +3300,9 @@ mod tests { &self, _request: Request, ) -> Result, Status> { - Ok(tonic::Response::new(CreateSandboxResponse {})) + Ok(tonic::Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( @@ -3228,6 +3380,8 @@ mod tests { sandbox_watch_bus: SandboxWatchBus::new(), tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), + supervisor_bootstrap_identity: None, + sandbox_claim_activation: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses: Vec::new(), diff --git a/crates/openshell-server/src/config_file.rs b/crates/openshell-server/src/config_file.rs index c4e0cbc959..3271c86881 100644 --- a/crates/openshell-server/src/config_file.rs +++ b/crates/openshell-server/src/config_file.rs @@ -129,7 +129,7 @@ pub struct GatewayFileSection { #[serde(default)] pub enable_user_namespaces: Option, /// Lifetime (seconds) of the projected `ServiceAccount` token kubelet - /// writes for the `IssueSandboxToken` bootstrap exchange. Driver + /// writes for the `RegisterSupervisorPod` bootstrap stream. Driver /// clamps to `[600, 86400]`. #[serde(default)] pub sa_token_ttl_secs: Option, diff --git a/crates/openshell-server/src/grpc/auth_rpc.rs b/crates/openshell-server/src/grpc/auth_rpc.rs index f4cb6a872a..84c0a47a66 100644 --- a/crates/openshell-server/src/grpc/auth_rpc.rs +++ b/crates/openshell-server/src/grpc/auth_rpc.rs @@ -3,8 +3,9 @@ //! Authentication-related RPC handlers. //! -//! Hosts the two sandbox-identity RPCs: -//! - `IssueSandboxToken` — bootstrap exchange (K8s SA token → gateway JWT) +//! Hosts the sandbox-identity RPCs: +//! - `RegisterSupervisorPod` — Kubernetes pod registration and activation +//! - `IssueSandboxToken` — legacy bootstrap compatibility shim //! - `RefreshSandboxToken` — renew a still-valid gateway JWT //! //! Both end in a fresh gateway-signed JWT minted by @@ -12,67 +13,60 @@ //! until their own `exp` and are bounded by the configured short TTL. use crate::ServerState; -use crate::auth::principal::{Principal, SandboxIdentitySource}; +use crate::auth::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; +use crate::warm_pod_activation::{load_sandbox, mint_pod_activation}; use openshell_core::proto::{ - IssueSandboxTokenRequest, IssueSandboxTokenResponse, RefreshSandboxTokenRequest, - RefreshSandboxTokenResponse, Sandbox, + IssueSandboxTokenRequest, IssueSandboxTokenResponse, PodActivationMessage, + RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, }; +use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentity; use std::sync::Arc; +use std::{pin::Pin, result::Result as StdResult}; +use tokio_stream::Stream; use tonic::{Request, Response, Status}; use tracing::{debug, info, warn}; +pub type RegisterSupervisorPodStream = + Pin> + Send + 'static>>; + #[allow(clippy::result_large_err, clippy::unused_async)] pub async fn handle_issue_sandbox_token( state: &Arc, request: Request, ) -> Result, Status> { - let principal = request - .extensions() - .get::() - .cloned() - .ok_or_else(|| Status::unauthenticated("missing principal"))?; - - let Principal::Sandbox(sandbox) = principal else { - return Err(Status::permission_denied( - "IssueSandboxToken requires a sandbox principal", - )); - }; + // Compatibility shim for older supervisor images. New Kubernetes + // supervisors should use RegisterSupervisorPod so warm-pool activation can + // later remain pending on the same bootstrap stream. + let sandbox = require_k8s_bootstrap_sandbox(request.extensions(), "IssueSandboxToken")?; + let activation = mint_pod_activation(state, &sandbox.sandbox_id, "IssueSandboxToken").await?; + Ok(Response::new(IssueSandboxTokenResponse { + token: activation.token, + expires_at_ms: activation.token_expires_at_ms, + })) +} - // Only the bootstrap K8s ServiceAccount path can mint a fresh gateway JWT - // via this RPC. Sandboxes already holding a gateway JWT use - // `RefreshSandboxToken` instead. - if !matches!( - sandbox.source, - SandboxIdentitySource::K8sServiceAccount { .. } - ) { - debug!( - sandbox_id = %sandbox.sandbox_id, - "IssueSandboxToken rejected: non-bootstrap principal source" +#[allow(clippy::result_large_err, clippy::unused_async)] +pub async fn handle_register_supervisor_pod( + state: &Arc, + request: Request, +) -> Result, Status> { + let identity = require_bootstrap_identity(request.extensions())?; + if let Some(sandbox_id) = identity.bound_sandbox_id() { + let activation = mint_pod_activation(state, sandbox_id, "RegisterSupervisorPod").await?; + info!( + sandbox_id = %activation.sandbox_id, + driver = %identity.driver, + instance_name = %identity.instance_name, + instance_id = %identity.instance_id, + "activated bound supervisor instance" ); - return Err(Status::permission_denied( - "this principal cannot mint a sandbox token; use RefreshSandboxToken", - )); + return Ok(Response::new(Box::pin(tokio_stream::once(Ok(activation))))); } - let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { - warn!( - sandbox_id = %sandbox.sandbox_id, - "IssueSandboxToken called but sandbox JWT issuer is not configured" - ); - Status::unavailable("sandbox JWT minting is not configured on this gateway") - })?; - - ensure_sandbox_exists(state, &sandbox.sandbox_id).await?; - - let minted = issuer.mint(&sandbox.sandbox_id)?; - info!( - sandbox_id = %sandbox.sandbox_id, - "issued gateway sandbox JWT" - ); - Ok(Response::new(IssueSandboxTokenResponse { - token: minted.token, - expires_at_ms: minted.expires_at_ms, - })) + let stream = state + .supervisor_pod_registrations + .register_pending(identity)?; + Ok(Response::new(Box::pin(stream))) } #[allow(clippy::result_large_err, clippy::unused_async)] @@ -92,15 +86,15 @@ pub async fn handle_refresh_sandbox_token( )); }; - // Only callers already holding a gateway-minted JWT may refresh; the - // K8s bootstrap path must use `IssueSandboxToken`. + // Only callers already holding a gateway-minted JWT may refresh; the K8s + // bootstrap path must use RegisterSupervisorPod. let SandboxIdentitySource::BootstrapJwt { .. } = &sandbox.source else { debug!( sandbox_id = %sandbox.sandbox_id, "RefreshSandboxToken rejected: non-gateway-JWT principal source" ); return Err(Status::permission_denied( - "this principal cannot refresh; use IssueSandboxToken for bootstrap", + "this principal cannot refresh; use RegisterSupervisorPod for bootstrap", )); }; @@ -112,7 +106,7 @@ pub async fn handle_refresh_sandbox_token( Status::unavailable("sandbox JWT minting is not configured on this gateway") })?; - ensure_sandbox_exists(state, &sandbox.sandbox_id).await?; + load_sandbox(state, &sandbox.sandbox_id).await?; let minted = issuer.mint(&sandbox.sandbox_id)?; info!( @@ -126,19 +120,57 @@ pub async fn handle_refresh_sandbox_token( })) } -async fn ensure_sandbox_exists(state: &Arc, sandbox_id: &str) -> Result<(), Status> { - if sandbox_id.is_empty() { - return Err(Status::invalid_argument("sandbox_id is required")); +fn require_k8s_bootstrap_sandbox( + extensions: &tonic::Extensions, + rpc_name: &'static str, +) -> Result { + let principal = extensions + .get::() + .cloned() + .ok_or_else(|| Status::unauthenticated("missing principal"))?; + + let Principal::Sandbox(sandbox) = principal else { + return Err(Status::permission_denied(format!( + "{rpc_name} requires a sandbox principal" + ))); + }; + + if !matches!( + sandbox.source, + SandboxIdentitySource::SupervisorBootstrap { .. } + ) { + debug!( + sandbox_id = %sandbox.sandbox_id, + rpc = rpc_name, + "bootstrap RPC rejected: non-bootstrap principal source" + ); + return Err(Status::permission_denied( + "this principal cannot mint a sandbox token; use RefreshSandboxToken", + )); } - state - .store - .get_message::(sandbox_id) - .await - .map_err(|e| Status::internal(format!("fetch sandbox failed: {e}")))? - .ok_or_else(|| Status::not_found("sandbox not found"))?; + Ok(sandbox) +} + +fn require_bootstrap_identity( + extensions: &tonic::Extensions, +) -> Result { + if let Some(identity) = extensions.get::().cloned() { + return Ok(identity); + } + + let principal = extensions + .get::() + .cloned() + .ok_or_else(|| Status::unauthenticated("missing principal"))?; + + let Principal::SupervisorBootstrap(identity) = principal else { + return Err(Status::permission_denied( + "RegisterSupervisorPod requires a supervisor bootstrap principal", + )); + }; - Ok(()) + Ok(identity) } #[cfg(test)] @@ -157,8 +189,10 @@ mod tests { use openshell_core::Config; use openshell_core::proto::datamodel::v1::ObjectMeta; use openshell_core::proto::{Sandbox, SandboxPhase, SandboxSpec}; + use openshell_core::supervisor_bootstrap::SupervisorBootstrapBinding; use std::collections::HashMap; use std::time::Duration; + use tokio_stream::StreamExt; async fn state_with_issuer() -> Arc { let mat = generate_jwt_key().expect("jwt key"); @@ -225,6 +259,17 @@ mod tests { }) } + fn bootstrap_identity(binding: SupervisorBootstrapBinding) -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "owner-uid-a".to_string(), + binding, + } + } + #[tokio::test] async fn refresh_returns_new_token() { let state = state_with_issuer().await; @@ -259,9 +304,10 @@ mod tests { req.extensions_mut() .insert(Principal::Sandbox(SandboxPrincipal { sandbox_id: "sandbox-a".to_string(), - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: "pod-a".to_string(), - pod_uid: "uid-a".to_string(), + source: SandboxIdentitySource::SupervisorBootstrap { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "uid-a".to_string(), }, trust_domain: Some("openshell".to_string()), })); @@ -273,6 +319,58 @@ mod tests { assert!(resp.expires_at_ms > 0); } + #[tokio::test] + async fn register_supervisor_pod_returns_immediate_activation_for_existing_sandbox() { + let state = state_with_issuer().await; + let mut req = Request::new(RegisterSupervisorPodRequest {}); + req.extensions_mut() + .insert(Principal::SupervisorBootstrap(bootstrap_identity( + SupervisorBootstrapBinding::BoundSandbox { + sandbox_id: "sandbox-a".to_string(), + }, + ))); + + let mut stream = handle_register_supervisor_pod(&state, req) + .await + .expect("register OK") + .into_inner(); + let activation = stream + .next() + .await + .expect("activation message") + .expect("activation OK"); + assert_eq!(activation.sandbox_id, "sandbox-a"); + assert_eq!(activation.sandbox_name, "sandbox-a"); + assert!(!activation.token.is_empty()); + assert!(activation.token_expires_at_ms > 0); + assert!(activation.startup_metadata.is_empty()); + assert!(stream.next().await.is_none()); + } + + #[tokio::test] + async fn register_supervisor_pod_keeps_unbound_warm_pod_pending() { + let state = state_with_issuer().await; + let mut req = Request::new(RegisterSupervisorPodRequest {}); + req.extensions_mut() + .insert(Principal::SupervisorBootstrap(bootstrap_identity( + SupervisorBootstrapBinding::WarmPending, + ))); + + let mut stream = handle_register_supervisor_pod(&state, req) + .await + .expect("register OK") + .into_inner(); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 1); + assert!( + tokio::time::timeout(Duration::from_millis(20), stream.next()) + .await + .is_err(), + "unbound warm pod must wait for later activation" + ); + drop(stream); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + } + #[tokio::test] async fn issue_rejects_missing_sandbox() { use crate::auth::principal::SandboxIdentitySource; @@ -282,9 +380,10 @@ mod tests { req.extensions_mut() .insert(Principal::Sandbox(SandboxPrincipal { sandbox_id: "sandbox-deleted".to_string(), - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: "pod-a".to_string(), - pod_uid: "uid-a".to_string(), + source: SandboxIdentitySource::SupervisorBootstrap { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "uid-a".to_string(), }, trust_domain: Some("openshell".to_string()), })); @@ -315,25 +414,26 @@ mod tests { } #[tokio::test] - async fn refresh_rejects_k8s_sa_principal() { - // K8s SA-bootstrap principals must use IssueSandboxToken, not - // RefreshSandboxToken — the refresh path assumes a still-valid - // gateway-minted JWT exists. + async fn refresh_rejects_bootstrap_principal() { + // Bootstrap principals must use RegisterSupervisorPod, not + // RefreshSandboxToken. The refresh path assumes a still-valid + // gateway-minted JWT already exists. use crate::auth::principal::SandboxIdentitySource; let state = state_with_issuer().await; let mut req = Request::new(RefreshSandboxTokenRequest {}); req.extensions_mut() .insert(Principal::Sandbox(SandboxPrincipal { sandbox_id: "sandbox-a".to_string(), - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: "pod-a".to_string(), - pod_uid: "uid-a".to_string(), + source: SandboxIdentitySource::SupervisorBootstrap { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "uid-a".to_string(), }, trust_domain: Some("openshell".to_string()), })); let err = handle_refresh_sandbox_token(&state, req) .await - .expect_err("K8s SA principal must not refresh"); + .expect_err("bootstrap principal must not refresh"); assert_eq!(err.code(), tonic::Code::PermissionDenied); } diff --git a/crates/openshell-server/src/grpc/mod.rs b/crates/openshell-server/src/grpc/mod.rs index ecd96ea90d..23331fbcf0 100644 --- a/crates/openshell-server/src/grpc/mod.rs +++ b/crates/openshell-server/src/grpc/mod.rs @@ -39,16 +39,17 @@ use openshell_core::proto::{ ListSandboxProvidersRequest, ListSandboxProvidersResponse, ListSandboxesRequest, ListSandboxesResponse, ListServicesRequest, ListServicesResponse, ListWorkspaceMembersRequest, ListWorkspaceMembersResponse, ListWorkspacesRequest, ListWorkspacesResponse, - ProviderProfileResponse, ProviderResponse, PushSandboxLogsRequest, PushSandboxLogsResponse, - RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RejectDraftChunkRequest, - RejectDraftChunkResponse, RelayFrame, RemoveWorkspaceMemberRequest, - RemoveWorkspaceMemberResponse, ReportPolicyStatusRequest, ReportPolicyStatusResponse, - RevokeSshSessionRequest, RevokeSshSessionResponse, RotateProviderCredentialRequest, - RotateProviderCredentialResponse, SandboxResponse, SandboxStreamEvent, ServiceEndpointResponse, - ServiceStatus, SubmitPolicyAnalysisRequest, SubmitPolicyAnalysisResponse, SupervisorMessage, - TcpForwardFrame, UndoDraftChunkRequest, UndoDraftChunkResponse, UpdateConfigRequest, - UpdateConfigResponse, UpdateProviderProfilesRequest, UpdateProviderProfilesResponse, - UpdateProviderRequest, WatchSandboxRequest, open_shell_server::OpenShell, + PodActivationMessage, ProviderProfileResponse, ProviderResponse, PushSandboxLogsRequest, + PushSandboxLogsResponse, RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, + RegisterSupervisorPodRequest, RejectDraftChunkRequest, RejectDraftChunkResponse, RelayFrame, + RemoveWorkspaceMemberRequest, RemoveWorkspaceMemberResponse, ReportPolicyStatusRequest, + ReportPolicyStatusResponse, RevokeSshSessionRequest, RevokeSshSessionResponse, + RotateProviderCredentialRequest, RotateProviderCredentialResponse, SandboxResponse, + SandboxStreamEvent, ServiceEndpointResponse, ServiceStatus, SubmitPolicyAnalysisRequest, + SubmitPolicyAnalysisResponse, SupervisorMessage, TcpForwardFrame, UndoDraftChunkRequest, + UndoDraftChunkResponse, UpdateConfigRequest, UpdateConfigResponse, + UpdateProviderProfilesRequest, UpdateProviderProfilesResponse, UpdateProviderRequest, + WatchSandboxRequest, open_shell_server::OpenShell, }; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; @@ -687,6 +688,18 @@ impl OpenShell for OpenShellService { auth_rpc::handle_issue_sandbox_token(&self.state, request).await } + type RegisterSupervisorPodStream = Pin< + Box> + Send + 'static>, + >; + + #[rpc_auth(auth = "pod_registration")] + async fn register_supervisor_pod( + &self, + request: Request, + ) -> Result, Status> { + auth_rpc::handle_register_supervisor_pod(&self.state, request).await + } + #[rpc_auth(auth = "sandbox")] async fn refresh_sandbox_token( &self, diff --git a/crates/openshell-server/src/grpc/policy.rs b/crates/openshell-server/src/grpc/policy.rs index 0703771846..ee1cc51eba 100644 --- a/crates/openshell-server/src/grpc/policy.rs +++ b/crates/openshell-server/src/grpc/policy.rs @@ -1252,6 +1252,9 @@ async fn resolve_sandbox_by_name_for_principal( Ok(sandbox) } Principal::User(_) => sandbox.ok_or_else(|| Status::not_found("sandbox not found")), + Principal::SupervisorBootstrap(_) => Err(Status::permission_denied( + "sandbox-scoped methods require a sandbox principal", + )), Principal::Anonymous => Err(Status::unauthenticated( "sandbox-scoped methods require an authenticated caller", )), diff --git a/crates/openshell-server/src/grpc/sandbox.rs b/crates/openshell-server/src/grpc/sandbox.rs index 57bf421b88..dfab3ad9f8 100644 --- a/crates/openshell-server/src/grpc/sandbox.rs +++ b/crates/openshell-server/src/grpc/sandbox.rs @@ -219,7 +219,7 @@ async fn handle_create_sandbox_inner( })?; // Mint the gateway JWT for singleplayer drivers. K8s sandboxes skip - // this mint and bootstrap via `IssueSandboxToken` at supervisor + // this mint and bootstrap via `RegisterSupervisorPod` at supervisor // startup; identifying "is this K8s?" lives in the compute layer, so // we mint unconditionally here when the issuer is configured and let // the K8s driver simply ignore the field. diff --git a/crates/openshell-server/src/inference.rs b/crates/openshell-server/src/inference.rs index 97217c2565..26d40d52c3 100644 --- a/crates/openshell-server/src/inference.rs +++ b/crates/openshell-server/src/inference.rs @@ -930,7 +930,10 @@ fn authorize_inference_bundle( ) -> Result { match principal { Some(crate::auth::principal::Principal::Sandbox(s)) => Ok(s.sandbox_id.clone()), - Some(crate::auth::principal::Principal::User(_)) => Err(Status::permission_denied( + Some( + crate::auth::principal::Principal::User(_) + | crate::auth::principal::Principal::SupervisorBootstrap(_), + ) => Err(Status::permission_denied( "GetInferenceBundle requires a sandbox principal", )), Some(crate::auth::principal::Principal::Anonymous) | None => Err(Status::unauthenticated( diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index a75534d397..24896d3c98 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -43,6 +43,7 @@ mod sandbox_index; mod sandbox_watch; mod service_routing; mod ssh_sessions; +mod supervisor_pod_registration; pub mod supervisor_session; mod telemetry; #[cfg(any(test, feature = "test-support"))] @@ -51,6 +52,7 @@ mod tls; #[cfg(test)] pub(crate) mod tls_test_utils; pub mod tracing_bus; +pub(crate) mod warm_pod_activation; mod ws_tunnel; use metrics_exporter_prometheus::PrometheusBuilder; @@ -132,11 +134,16 @@ pub struct ServerState { /// Validated built-in and operator-registered supervisor middleware. pub middleware_registry: Arc, + /// Pending Kubernetes supervisor pod registrations awaiting warm-pool + /// activation. + pub(crate) supervisor_pod_registrations: + Arc, + /// OIDC JWKS cache for JWT validation. `None` when OIDC is not configured. pub oidc_cache: Option>, /// Gateway-minted sandbox JWT issuer. `None` when `config.gateway_jwt` - /// is not configured; in that mode `IssueSandboxToken` returns + /// is not configured; in that mode Kubernetes bootstrap RPCs return /// `Status::unavailable`. Populated at startup from the on-disk key /// material that `certgen` writes. pub sandbox_jwt_issuer: Option>, @@ -146,11 +153,6 @@ pub struct ServerState { /// presenting a freshly minted token are recognized. pub sandbox_jwt_authenticator: Option>, - /// Optional K8s `ServiceAccount` authenticator that backs the - /// `IssueSandboxToken` bootstrap path. Only present when the gateway - /// runs in-cluster. - pub k8s_sa_authenticator: Option>, - /// Gateway-wide gRPC request rate limiter shared by every multiplex path. pub(crate) grpc_rate_limiter: Option, @@ -207,10 +209,12 @@ impl ServerState { settings_mutex: tokio::sync::Mutex::new(()), supervisor_sessions, middleware_registry: Arc::new(MiddlewareRegistry::default()), + supervisor_pod_registrations: Arc::new( + supervisor_pod_registration::SupervisorPodRegistrationRegistry::new(), + ), oidc_cache, sandbox_jwt_issuer: None, sandbox_jwt_authenticator: None, - k8s_sa_authenticator: None, grpc_rate_limiter, gateway_interceptors: None, provider_profile_sources: @@ -385,42 +389,6 @@ pub(crate) async fn run_server( state.sandbox_jwt_authenticator = Some(Arc::new(authenticator)); } - // K8s ServiceAccount bootstrap authenticator. Only constructed when - // the gateway is running in-cluster (kubelet provides the API host - // env var) and has a sandbox JWT issuer to mint replacements against; - // outside the cluster we can't call the apiserver's TokenReview API, - // and without the issuer there's nothing to exchange the SA token for. - if state.sandbox_jwt_issuer.is_some() && std::env::var_os("KUBERNETES_SERVICE_HOST").is_some() { - // Pod lookups and TokenReview identity checks must match the sandbox - // namespace and service account used by the Kubernetes driver. - let kubernetes_config = - compute::driver_config::kubernetes_config_for_k8s_sa_bootstrap(config_file.as_ref())?; - let sandbox_namespace = kubernetes_config.namespace; - let sandbox_service_account = kubernetes_config.service_account_name; - match kube::Client::try_default().await { - Ok(client) => { - let resolver = Arc::new(auth::k8s_sa::LiveK8sResolver::new( - client, - &sandbox_namespace, - "openshell-gateway".to_string(), - sandbox_service_account.clone(), - )); - let authenticator = auth::k8s_sa::K8sServiceAccountAuthenticator::new(resolver); - state.k8s_sa_authenticator = Some(Arc::new(authenticator)); - info!( - namespace = %sandbox_namespace, - service_account = %sandbox_service_account, - "K8s ServiceAccount bootstrap authenticator enabled" - ); - } - Err(e) => warn!( - error = %e, - "in-cluster K8s client construction failed; \ - K8s ServiceAccount bootstrap is disabled" - ), - } - } - let state = Arc::new(state); let (shutdown_tx, shutdown_rx) = watch::channel(false); @@ -436,6 +404,10 @@ pub(crate) async fn run_server( } state.compute.spawn_watchers(shutdown_rx.clone()); + state.compute.spawn_sandbox_claim_activation( + Arc::new(warm_pod_activation::GatewaySupervisorBootstrapActivator::new(state.clone())), + shutdown_rx.clone(), + ); ssh_sessions::spawn_session_reaper(store.clone(), Duration::from_secs(3600)); supervisor_session::spawn_relay_reaper(state.clone(), Duration::from_secs(30)); provider_refresh::spawn_refresh_worker(state.clone(), Duration::from_secs(60)); diff --git a/crates/openshell-server/src/multiplex.rs b/crates/openshell-server/src/multiplex.rs index 83e27d27cb..65d5d349e3 100644 --- a/crates/openshell-server/src/multiplex.rs +++ b/crates/openshell-server/src/multiplex.rs @@ -463,7 +463,7 @@ fn gateway_principal_fields(principal: &Principal) -> BTreeMap { match &sandbox.source { SandboxIdentitySource::BootstrapJwt { .. } => "bootstrap_jwt", SandboxIdentitySource::BootstrapCert { .. } => "bootstrap_cert", - SandboxIdentitySource::K8sServiceAccount { .. } => "k8s_service_account", + SandboxIdentitySource::SupervisorBootstrap { .. } => "supervisor_bootstrap", } .to_string(), ); @@ -471,6 +471,14 @@ fn gateway_principal_fields(principal: &Principal) -> BTreeMap { fields.insert("trust_domain".to_string(), trust_domain.clone()); } } + Principal::SupervisorBootstrap(identity) => { + fields.insert("kind".to_string(), "supervisor_bootstrap".to_string()); + fields.insert("driver".to_string(), identity.driver.clone()); + fields.insert("instance_id".to_string(), identity.instance_id.clone()); + fields.insert("instance_name".to_string(), identity.instance_name.clone()); + fields.insert("owner_name".to_string(), identity.owner_name.clone()); + fields.insert("owner_uid".to_string(), identity.owner_uid.clone()); + } Principal::Anonymous => { fields.insert("kind".to_string(), "anonymous".to_string()); } @@ -694,10 +702,12 @@ where /// Assemble the authenticator chain for the gateway. /// /// Chain order (first-match-wins): -/// 1. `K8sServiceAccountAuthenticator` (path-scoped to `IssueSandboxToken`) -/// — exchanges a projected SA token for a `Principal::Sandbox` so the -/// `IssueSandboxToken` handler can mint a gateway JWT. No-op on every -/// other path; only present when the gateway runs in-cluster. +/// 1. `SupervisorBootstrapAuthenticator` (path-scoped to supervisor bootstrap +/// RPCs) — delegates driver-native bootstrap token validation to the active +/// compute driver and resolves it to either a legacy `Principal::Sandbox` +/// for `IssueSandboxToken` or a registration-scoped principal for +/// `RegisterSupervisorPod`. No-op on every other path; only present when +/// the active driver exposes a bootstrap identity provider. /// 2. `SandboxJwtAuthenticator` — validates gateway-minted JWTs. Recognized /// via a distinctive `kid` so non-matching Bearer tokens fall through. /// 3. `OidcAuthenticator` — validates user Bearer tokens against the @@ -715,8 +725,10 @@ where /// to pass-through unless mTLS or local unauthenticated users are enabled. fn build_authenticator_chain(state: &ServerState) -> Option { let mut authenticators: Vec> = Vec::new(); - if let Some(k8s) = state.k8s_sa_authenticator.clone() { - authenticators.push(k8s); + if let Some(provider) = state.compute.supervisor_bootstrap_identity_provider() { + authenticators.push(Arc::new( + crate::auth::k8s_sa::SupervisorBootstrapAuthenticator::new(provider), + )); } if let Some(jwt) = state.sandbox_jwt_authenticator.clone() { authenticators.push(jwt); @@ -743,6 +755,8 @@ fn build_authenticator_chain(state: &ServerState) -> Option /// `Principal::User` is gated by the RBAC `AuthzPolicy`. /// `Principal::Sandbox` is gated by a supervisor-method allowlist, then /// handlers enforce same-sandbox scope on request bodies. +/// `Principal::SupervisorBootstrap` is gated by the pod-registration method +/// only. #[derive(Clone)] pub struct AuthGrpcRouter { inner: S, @@ -881,6 +895,14 @@ where ))); } } + Principal::SupervisorBootstrap(ref identity) => { + if !crate::auth::method_authz::is_pod_registration_callable(&path) { + return Ok(status_response(tonic::Status::permission_denied( + "supervisor bootstrap principals may only register supervisors", + ))); + } + req.extensions_mut().insert(identity.clone()); + } Principal::Anonymous => { return Ok(status_response(tonic::Status::unauthenticated( "anonymous callers may not call authenticated methods", @@ -1819,6 +1841,9 @@ mod tests { Principal, SandboxIdentitySource, SandboxPrincipal, UserPrincipal, }; use http_body_util::Full; + use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapBinding, SupervisorBootstrapIdentity, + }; use std::sync::Arc; use std::sync::Mutex; use tower::Service; @@ -1906,6 +1931,17 @@ mod tests { }) } + fn bootstrap_registration_principal() -> Principal { + Principal::SupervisorBootstrap(SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "pod-a".to_string(), + instance_id: "pod-uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "owner-uid-a".to_string(), + binding: SupervisorBootstrapBinding::WarmPending, + }) + } + #[tokio::test] async fn mtls_peer_identity_fills_missing_principal_when_enabled() { let mock = Arc::new(MockAuthenticator::returning(Ok(None))); @@ -2094,6 +2130,43 @@ mod tests { )); } + #[tokio::test] + async fn bootstrap_registration_principal_can_only_register_supervisor_pod() { + let mock = Arc::new(MockAuthenticator::returning(Ok(Some( + bootstrap_registration_principal(), + )))); + let chain = AuthenticatorChain::new(vec![mock]); + let (recorder, seen) = PrincipalRecorder::new(); + let mut router = AuthGrpcRouter::new(recorder, Some(chain), None); + + let res = router + .call(empty_request( + "/openshell.v1.OpenShell/RegisterSupervisorPod", + )) + .await + .unwrap(); + + assert_eq!(res.status(), 200); + assert!(matches!( + seen.lock().unwrap().as_ref(), + Some(Principal::SupervisorBootstrap(_)) + )); + + let mock = Arc::new(MockAuthenticator::returning(Ok(Some( + bootstrap_registration_principal(), + )))); + let chain = AuthenticatorChain::new(vec![mock]); + let (recorder, seen) = PrincipalRecorder::new(); + let mut router = AuthGrpcRouter::new(recorder, Some(chain), None); + let res = router + .call(empty_request("/openshell.v1.OpenShell/GetSandboxConfig")) + .await + .unwrap(); + + assert!(seen.lock().unwrap().is_none()); + assert_eq!(grpc_status(&res).as_deref(), Some("7")); + } + /// A user principal — even one carrying `openshell:all` and the /// admin role — must not reach a `sandbox`-annotated method. The /// router enforces this from the per-handler auth-mode declarations @@ -2150,6 +2223,7 @@ mod tests { "/openshell.v1.OpenShell/CreateProvider", "/openshell.v1.OpenShell/ApproveDraftChunk", "/openshell.inference.v1.Inference/GetInferenceRoute", + "/openshell.v1.OpenShell/RegisterSupervisorPod", "/openshell.inference.v1.Inference/SetInferenceRoute", ] { let mock = Arc::new(MockAuthenticator::returning(Ok(Some(sandbox_principal())))); diff --git a/crates/openshell-server/src/persistence/mod.rs b/crates/openshell-server/src/persistence/mod.rs index 291e2eafd7..d8e6267c25 100644 --- a/crates/openshell-server/src/persistence/mod.rs +++ b/crates/openshell-server/src/persistence/mod.rs @@ -13,6 +13,7 @@ pub use openshell_core::proto::{ use openshell_core::{Error as CoreError, Result as CoreResult}; use prost::Message; use rand::Rng; +use sha2::{Digest, Sha256}; use std::collections::HashMap; use thiserror::Error; @@ -23,6 +24,12 @@ pub use sqlite::SqliteStore; pub const POLICY_OBJECT_TYPE: &str = "sandbox_policy"; /// Object type string for draft policy chunk records. pub const DRAFT_CHUNK_OBJECT_TYPE: &str = "draft_policy_chunk"; +/// Object type string for warm-pool activation token hashes. +pub const ACTIVATION_TOKEN_OBJECT_TYPE: &str = "activation_token"; + +fn activation_token_record_id(sandbox_id: &str) -> String { + format!("{ACTIVATION_TOKEN_OBJECT_TYPE}/{sandbox_id}") +} pub type PersistenceResult = Result; @@ -210,6 +217,88 @@ impl Store { store_dispatch!(self.ping()) } + /// Hash an opaque warm-pool activation token for durable storage. + #[must_use] + pub fn activation_token_hash(token: &str) -> String { + let digest = Sha256::digest(token.as_bytes()); + hex::encode(digest) + } + + /// Store the hash of a driver-issued activation token for a sandbox. + /// + /// Re-inserting the same token for the same sandbox is idempotent. Changing + /// the token for a sandbox, or binding an existing token hash to another + /// sandbox, is rejected. + pub async fn put_activation_token( + &self, + sandbox_id: &str, + token: &str, + ) -> PersistenceResult<()> { + if sandbox_id.is_empty() { + return Err(PersistenceError::Encode( + "sandbox_id is required for activation token mapping".to_string(), + )); + } + if token.is_empty() { + return Err(PersistenceError::Encode( + "activation token cannot be empty".to_string(), + )); + } + + let token_hash = Self::activation_token_hash(token); + if let Some(existing) = self.get_activation_token_hash(sandbox_id).await? { + if existing == token_hash { + return Ok(()); + } + return Err(PersistenceError::Conflict { + current_resource_version: None, + }); + } + + let payload = token_hash.as_bytes(); + self.put_if( + ACTIVATION_TOKEN_OBJECT_TYPE, + &activation_token_record_id(sandbox_id), + &token_hash, + "", + payload, + None, + WriteCondition::MustCreate, + ) + .await + .map(|_| ()) + } + + /// Fetch a stored activation token hash for a sandbox id. + pub async fn get_activation_token_hash( + &self, + sandbox_id: &str, + ) -> PersistenceResult> { + self.get( + ACTIVATION_TOKEN_OBJECT_TYPE, + &activation_token_record_id(sandbox_id), + ) + .await? + .map(|record| { + String::from_utf8(record.payload).map_err(|e| { + PersistenceError::Decode(format!("activation token hash is not valid UTF-8: {e}")) + }) + }) + .transpose() + } + + /// Delete a sandbox's activation token mapping. + pub async fn delete_activation_token_for_sandbox( + &self, + sandbox_id: &str, + ) -> PersistenceResult { + self.delete( + ACTIVATION_TOKEN_OBJECT_TYPE, + &activation_token_record_id(sandbox_id), + ) + .await + } + /// Test support only: close the underlying connection pool. /// /// There is no runtime shutdown path yet. If we add graceful shutdown, diff --git a/crates/openshell-server/src/persistence/tests.rs b/crates/openshell-server/src/persistence/tests.rs index 9539eab49d..de7716655b 100644 --- a/crates/openshell-server/src/persistence/tests.rs +++ b/crates/openshell-server/src/persistence/tests.rs @@ -1,7 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use super::{ObjectType, PersistenceError, Store, generate_name, test_store}; +use super::{ + ACTIVATION_TOKEN_OBJECT_TYPE, ObjectType, PersistenceError, Store, generate_name, test_store, +}; use crate::policy_store::{AtomicPolicyRevisionWrite, PolicyStoreExt}; use openshell_core::proto::datamodel::v1::ObjectMeta as ProtoObjectMeta; use openshell_core::proto::{ObjectForTest, Sandbox, SandboxPolicy, SandboxSpec}; @@ -266,6 +268,97 @@ async fn sqlite_delete_behavior() { assert!(!deleted_again); } +#[tokio::test] +async fn activation_token_mapping_stores_hash_and_deletes() { + let store = test_store().await; + + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + let hash = store + .get_activation_token_hash("sandbox-a") + .await + .unwrap() + .unwrap(); + + assert_eq!(hash, Store::activation_token_hash("raw-activation-token")); + assert_ne!(hash, "raw-activation-token"); + + let record = store + .get(ACTIVATION_TOKEN_OBJECT_TYPE, "activation_token/sandbox-a") + .await + .unwrap() + .unwrap(); + assert_eq!(record.payload, hash.as_bytes()); + assert_ne!(record.payload, b"raw-activation-token"); + + assert!( + store + .delete_activation_token_for_sandbox("sandbox-a") + .await + .unwrap() + ); + assert!( + store + .get_activation_token_hash("sandbox-a") + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn activation_token_mapping_is_idempotent_for_same_token() { + let store = test_store().await; + + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + + assert_eq!( + store.get_activation_token_hash("sandbox-a").await.unwrap(), + Some(Store::activation_token_hash("raw-activation-token")) + ); +} + +#[tokio::test] +async fn activation_token_mapping_rejects_token_change_for_sandbox() { + let store = test_store().await; + + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + let err = store + .put_activation_token("sandbox-a", "other-token") + .await + .unwrap_err(); + + assert!(matches!(err, PersistenceError::Conflict { .. })); +} + +#[tokio::test] +async fn activation_token_mapping_rejects_token_reuse_for_other_sandbox() { + let store = test_store().await; + + store + .put_activation_token("sandbox-a", "raw-activation-token") + .await + .unwrap(); + let err = store + .put_activation_token("sandbox-b", "raw-activation-token") + .await + .unwrap_err(); + + assert!(matches!(err, PersistenceError::UniqueViolation { .. })); +} + #[tokio::test] async fn sqlite_protobuf_round_trip() { let store = test_store().await; diff --git a/crates/openshell-server/src/supervisor_pod_registration.rs b/crates/openshell-server/src/supervisor_pod_registration.rs new file mode 100644 index 0000000000..805cfc27a2 --- /dev/null +++ b/crates/openshell-server/src/supervisor_pod_registration.rs @@ -0,0 +1,438 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Pending supervisor pod registrations for warm-pool activation. +//! +//! Cold pods already bound to a sandbox are activated immediately by the +//! `RegisterSupervisorPod` handler. Warm pods can register before claim +//! assignment; the future claim controller will activate the stored stream once +//! it binds the exact pod UID to a sandbox. + +use openshell_core::proto::PodActivationMessage; +use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentity; +use std::collections::HashMap; +use std::pin::Pin; +use std::result::Result as StdResult; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; +use std::time::Instant; +use tokio::sync::mpsc; +use tokio_stream::Stream; +use tokio_stream::wrappers::ReceiverStream; +use tonic::Status; +use tracing::{debug, info, warn}; + +#[derive(Debug, Default)] +pub struct SupervisorPodRegistrationRegistry { + inner: Mutex, + next_session_id: AtomicU64, +} + +#[derive(Debug, Default)] +struct Inner { + pending_by_instance_id: HashMap, + activated_instance_id: HashMap, +} + +#[derive(Debug)] +struct PendingRegistration { + identity: SupervisorBootstrapIdentity, + sender: mpsc::Sender>, + session_id: u64, + registered_at: Instant, +} + +impl SupervisorPodRegistrationRegistry { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + #[allow(clippy::result_large_err)] + pub fn register_pending( + self: &Arc, + identity: SupervisorBootstrapIdentity, + ) -> Result { + if identity.instance_id.is_empty() { + return Err(Status::permission_denied( + "registered supervisor instance ID is required", + )); + } + + let (sender, receiver) = mpsc::channel(1); + let session_id = self.next_session_id.fetch_add(1, Ordering::Relaxed); + let instance_id = identity.instance_id.clone(); + let instance_name = identity.instance_name.clone(); + let owner_name = identity.owner_name.clone(); + let driver = identity.driver.clone(); + + let replaced = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_instance_id.contains_key(&instance_id) { + return Err(Status::already_exists( + "supervisor instance has already been activated", + )); + } + inner + .pending_by_instance_id + .insert( + instance_id.clone(), + PendingRegistration { + identity, + sender, + session_id, + registered_at: Instant::now(), + }, + ) + .is_some() + }; + + if replaced { + info!( + driver = %driver, + instance = %instance_name, + instance_id = %instance_id, + owner = %owner_name, + "replaced duplicate pending supervisor registration" + ); + } else { + info!( + driver = %driver, + instance = %instance_name, + instance_id = %instance_id, + owner = %owner_name, + "registered warm supervisor instance pending activation" + ); + } + + Ok(PendingRegistrationStream { + registry: self.clone(), + instance_id, + session_id, + inner: ReceiverStream::new(receiver), + }) + } + + #[allow(clippy::result_large_err)] + pub(crate) fn pending_identity( + &self, + instance_id: &str, + ) -> Result { + let inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_instance_id.contains_key(instance_id) { + return Err(Status::already_exists( + "supervisor instance has already been activated", + )); + } + inner + .pending_by_instance_id + .get(instance_id) + .map(|pending| pending.identity.clone()) + .ok_or_else(|| Status::not_found("pending supervisor registration not found")) + } + + #[allow(clippy::result_large_err)] + pub(crate) fn activate( + &self, + instance_id: &str, + activation: PodActivationMessage, + ) -> Result<(), Status> { + let pending = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_instance_id.contains_key(instance_id) { + return Err(Status::already_exists( + "supervisor instance has already been activated", + )); + } + let Some(pending) = inner.pending_by_instance_id.remove(instance_id) else { + debug!( + instance_id = %instance_id, + "no pending supervisor registration to activate" + ); + return Err(Status::not_found( + "pending supervisor registration not found", + )); + }; + inner + .activated_instance_id + .insert(instance_id.to_string(), Instant::now()); + pending + }; + + let instance_name = pending.identity.instance_name.clone(); + if pending.sender.try_send(Ok(activation)).is_err() { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + inner.activated_instance_id.remove(instance_id); + warn!( + instance = %instance_name, + instance_id = %instance_id, + "pending supervisor registration stream closed before activation" + ); + return Err(Status::unavailable( + "registered supervisor stream closed before activation", + )); + } + + info!( + instance = %instance_name, + instance_id = %instance_id, + pending_ms = pending.registered_at.elapsed().as_millis(), + "activated pending supervisor registration" + ); + Ok(()) + } + + #[allow(clippy::result_large_err)] + pub(crate) fn fail_pending(&self, instance_id: &str, status: Status) -> Result<(), Status> { + let pending = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_instance_id.contains_key(instance_id) { + return Err(Status::already_exists( + "supervisor instance has already been activated", + )); + } + inner.pending_by_instance_id.remove(instance_id) + }; + + let Some(pending) = pending else { + debug!(instance_id = %instance_id, "no pending supervisor registration to fail"); + return Err(Status::not_found( + "pending supervisor registration not found", + )); + }; + + let instance_name = pending.identity.instance_name.clone(); + pending.sender.try_send(Err(status)).map_err(|_| { + warn!( + instance = %instance_name, + instance_id = %instance_id, + "pending supervisor registration stream closed before failure notification" + ); + Status::unavailable("registered supervisor stream closed before failure notification") + })?; + info!( + instance = %instance_name, + instance_id = %instance_id, + pending_ms = pending.registered_at.elapsed().as_millis(), + "failed pending supervisor registration" + ); + Ok(()) + } + + #[must_use] + #[allow(dead_code)] + pub fn pending_count(&self) -> usize { + self.inner + .lock() + .expect("pending pod registry poisoned") + .pending_by_instance_id + .len() + } + + #[must_use] + #[allow(dead_code)] + pub fn activated_count(&self) -> usize { + self.inner + .lock() + .expect("pending pod registry poisoned") + .activated_instance_id + .len() + } + + fn remove_if_session(&self, instance_id: &str, session_id: u64) { + let removed = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner + .pending_by_instance_id + .get(instance_id) + .is_some_and(|pending| pending.session_id == session_id) + { + inner.pending_by_instance_id.remove(instance_id) + } else { + None + } + }; + + if let Some(pending) = removed { + debug!( + instance = %pending.identity.instance_name, + instance_id = %instance_id, + pending_ms = pending.registered_at.elapsed().as_millis(), + "removed pending supervisor registration" + ); + } + } +} + +pub struct PendingRegistrationStream { + registry: Arc, + instance_id: String, + session_id: u64, + inner: ReceiverStream>, +} + +impl Stream for PendingRegistrationStream { + type Item = StdResult; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.inner).poll_next(cx) + } +} + +impl Drop for PendingRegistrationStream { + fn drop(&mut self) { + self.registry + .remove_if_session(&self.instance_id, self.session_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + use tokio_stream::StreamExt; + + fn bootstrap_identity(instance_id: &str) -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "warm-pod-a".to_string(), + instance_id: instance_id.to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "owner-uid-a".to_string(), + binding: openshell_core::supervisor_bootstrap::SupervisorBootstrapBinding::WarmPending, + } + } + + #[test] + fn dropping_stream_removes_pending_registration() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let stream = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register pending"); + assert_eq!(registry.pending_count(), 1); + drop(stream); + assert_eq!(registry.pending_count(), 0); + } + + #[test] + fn duplicate_registration_replaces_prior_stream() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let old = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register old"); + let new = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register new"); + + assert_eq!(registry.pending_count(), 1); + drop(old); + assert_eq!( + registry.pending_count(), + 1, + "old stream drop must not remove replacement" + ); + drop(new); + assert_eq!(registry.pending_count(), 0); + } + + #[tokio::test] + async fn activation_sends_message_and_removes_pending_registration() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let mut stream = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register pending"); + + let activation = PodActivationMessage { + sandbox_id: "sandbox-a".to_string(), + sandbox_name: "sandbox-a".to_string(), + token: "token-a".to_string(), + token_expires_at_ms: 123, + startup_metadata: HashMap::default(), + }; + registry + .activate("pod-uid-a", activation) + .expect("activate"); + assert_eq!(registry.pending_count(), 0); + assert_eq!(registry.activated_count(), 1); + + let received = stream + .next() + .await + .expect("activation message") + .expect("activation OK"); + assert_eq!(received.sandbox_id, "sandbox-a"); + assert!(stream.next().await.is_none()); + } + + #[test] + fn activation_for_unknown_pod_uid_returns_not_found() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let activation = PodActivationMessage { + sandbox_id: "sandbox-a".to_string(), + sandbox_name: "sandbox-a".to_string(), + token: "token-a".to_string(), + token_expires_at_ms: 123, + startup_metadata: HashMap::default(), + }; + + let err = registry + .activate("pod-uid-a", activation) + .expect_err("unknown pod UID must fail"); + + assert_eq!(err.code(), tonic::Code::NotFound); + } + + #[tokio::test] + async fn activation_tombstone_rejects_duplicate_activation_and_registration() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let mut stream = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register pending"); + + let activation = PodActivationMessage { + sandbox_id: "sandbox-a".to_string(), + sandbox_name: "sandbox-a".to_string(), + token: "token-a".to_string(), + token_expires_at_ms: 123, + startup_metadata: HashMap::default(), + }; + registry + .activate("pod-uid-a", activation.clone()) + .expect("activate"); + let _ = stream.next().await.expect("activation message"); + + let err = registry + .activate("pod-uid-a", activation) + .expect_err("duplicate activation must fail"); + assert_eq!(err.code(), tonic::Code::AlreadyExists); + + let Err(err) = registry.register_pending(bootstrap_identity("pod-uid-a")) else { + panic!("activated pod UID must not register again"); + }; + assert_eq!(err.code(), tonic::Code::AlreadyExists); + } + + #[test] + fn closed_pending_stream_cannot_be_activated() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let stream = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register pending"); + drop(stream); + + let activation = PodActivationMessage { + sandbox_id: "sandbox-a".to_string(), + sandbox_name: "sandbox-a".to_string(), + token: "token-a".to_string(), + token_expires_at_ms: 123, + startup_metadata: HashMap::default(), + }; + let err = registry + .activate("pod-uid-a", activation) + .expect_err("closed stream should remove pending registration"); + + assert_eq!(err.code(), tonic::Code::NotFound); + } +} diff --git a/crates/openshell-server/src/test_support.rs b/crates/openshell-server/src/test_support.rs index 9cd80d6ed8..589d55105c 100644 --- a/crates/openshell-server/src/test_support.rs +++ b/crates/openshell-server/src/test_support.rs @@ -257,7 +257,9 @@ impl ComputeDriver for FakeComputeDriver { .calls .push(FakeComputeDriverCall::CreateSandbox { sandbox }); }); - Ok(Response::new(CreateSandboxResponse {})) + Ok(Response::new(CreateSandboxResponse { + activation_token: String::new(), + })) } async fn stop_sandbox( diff --git a/crates/openshell-server/src/warm_pod_activation.rs b/crates/openshell-server/src/warm_pod_activation.rs new file mode 100644 index 0000000000..e19b6a135a --- /dev/null +++ b/crates/openshell-server/src/warm_pod_activation.rs @@ -0,0 +1,625 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Gateway-side warm supervisor pod activation. +//! +//! The Kubernetes claim controller will call this module after it observes and +//! revalidates a claim that binds a warm supervisor instance to an `OpenShell` +//! sandbox. + +use crate::ServerState; +use crate::persistence::Store; +use async_trait::async_trait; +use openshell_core::proto::{PodActivationMessage, Sandbox}; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapActivationRequest, SupervisorBootstrapActivator, SupervisorBootstrapIdentity, +}; +use std::sync::Arc; +use tonic::Status; +use tracing::{info, warn}; + +#[derive(Debug, Clone, PartialEq, Eq)] +#[allow(dead_code)] +pub struct WarmPodActivationTarget { + pub driver: String, + pub instance_id: String, + pub sandbox_id: String, + pub owner_uid: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +#[allow(dead_code)] +pub struct ValidatedWarmPodActivation { + pub driver: String, + pub instance_id: String, + pub instance_name: String, + pub owner_uid: String, + pub sandbox_id: String, +} + +#[derive(Debug)] +pub enum WarmPodActivationError { + RetryLater(Status), + Fatal(Status), +} + +impl WarmPodActivationError { + fn fatal(status: Status) -> Self { + Self::Fatal(status) + } + + fn retry_later(status: Status) -> Self { + Self::RetryLater(status) + } +} + +#[async_trait] +#[allow(dead_code)] +pub trait WarmPodActivationValidator: Send + Sync { + async fn validate( + &self, + target: &WarmPodActivationTarget, + pending: &SupervisorBootstrapIdentity, + ) -> Result; +} + +#[derive(Debug)] +#[allow(dead_code)] +pub struct GatewaySupervisorBootstrapActivator { + state: Arc, +} + +impl GatewaySupervisorBootstrapActivator { + #[must_use] + #[allow(dead_code)] + pub fn new(state: Arc) -> Self { + Self { state } + } +} + +#[async_trait] +impl SupervisorBootstrapActivator for GatewaySupervisorBootstrapActivator { + async fn activate_registered_supervisor( + &self, + request: SupervisorBootstrapActivationRequest, + ) -> Result<(), Status> { + let target = WarmPodActivationTarget { + driver: request.driver, + instance_id: request.instance_id, + sandbox_id: request.sandbox_id, + owner_uid: request.owner_uid, + }; + let validator = TokenBackedActivationRequest { + state: self.state.clone(), + activation_token: request.activation_token, + }; + activate_warm_pod(&self.state, &validator, target).await + } +} + +struct TokenBackedActivationRequest { + state: Arc, + activation_token: String, +} + +#[async_trait] +impl WarmPodActivationValidator for TokenBackedActivationRequest { + async fn validate( + &self, + target: &WarmPodActivationTarget, + pending: &SupervisorBootstrapIdentity, + ) -> Result { + if self.activation_token.is_empty() { + return Err(WarmPodActivationError::fatal(Status::invalid_argument( + "activation token is required", + ))); + } + + let expected_hash = self + .state + .store + .get_activation_token_hash(&target.sandbox_id) + .await + .map_err(|err| { + WarmPodActivationError::retry_later(Status::unavailable(format!( + "activation token mapping unavailable: {err}" + ))) + })? + .ok_or_else(|| { + WarmPodActivationError::retry_later(Status::unavailable( + "activation token mapping is not available yet", + )) + })?; + let request_hash = Store::activation_token_hash(&self.activation_token); + if request_hash != expected_hash { + return Err(WarmPodActivationError::fatal(Status::permission_denied( + "activation token does not match sandbox", + ))); + } + + Ok(ValidatedWarmPodActivation { + driver: target.driver.clone(), + instance_id: target.instance_id.clone(), + instance_name: pending.instance_name.clone(), + owner_uid: target.owner_uid.clone(), + sandbox_id: target.sandbox_id.clone(), + }) + } +} + +#[allow(clippy::result_large_err)] +#[allow(dead_code)] +pub async fn activate_warm_pod( + state: &Arc, + validator: &V, + target: WarmPodActivationTarget, +) -> Result<(), Status> +where + V: WarmPodActivationValidator + ?Sized, +{ + let pending = state + .supervisor_pod_registrations + .pending_identity(&target.instance_id)?; + let activation_result = async { + let validated = validator.validate(&target, &pending).await?; + ensure_validated_activation_matches_pending(&target, &pending, &validated) + .map_err(WarmPodActivationError::fatal)?; + mint_pod_activation(state, &validated.sandbox_id, "WarmPodActivation") + .await + .map_err(|status| { + if status.code() == tonic::Code::NotFound { + WarmPodActivationError::retry_later(status) + } else { + WarmPodActivationError::fatal(status) + } + }) + } + .await; + let activation = match activation_result { + Ok(activation) => activation, + Err(WarmPodActivationError::RetryLater(status)) => { + return Err(status); + } + Err(WarmPodActivationError::Fatal(status)) => { + let stream_status = clone_status(&status); + state + .supervisor_pod_registrations + .fail_pending(&target.instance_id, stream_status)?; + return Err(status); + } + }; + + state + .supervisor_pod_registrations + .activate(&target.instance_id, activation)?; + Ok(()) +} + +fn clone_status(status: &Status) -> Status { + Status::new(status.code(), status.message().to_string()) +} + +#[allow(clippy::result_large_err)] +fn ensure_validated_activation_matches_pending( + target: &WarmPodActivationTarget, + pending: &SupervisorBootstrapIdentity, + validated: &ValidatedWarmPodActivation, +) -> Result<(), Status> { + if validated.driver != target.driver || validated.driver != pending.driver { + return Err(Status::permission_denied( + "validated driver does not match pending registration", + )); + } + if validated.instance_id != target.instance_id || validated.instance_id != pending.instance_id { + return Err(Status::permission_denied( + "validated instance ID does not match pending registration", + )); + } + if validated.instance_name != pending.instance_name { + return Err(Status::permission_denied( + "validated instance name does not match pending registration", + )); + } + if validated.owner_uid != target.owner_uid || validated.owner_uid != pending.owner_uid { + return Err(Status::permission_denied( + "validated owner UID does not match pending registration", + )); + } + if validated.sandbox_id != target.sandbox_id { + return Err(Status::permission_denied( + "validated sandbox ID does not match activation target", + )); + } + Ok(()) +} + +#[allow(clippy::result_large_err)] +pub async fn mint_pod_activation( + state: &Arc, + sandbox_id: &str, + reason: &'static str, +) -> Result { + let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { + warn!( + sandbox_id = %sandbox_id, + reason, + "pod activation requested but sandbox JWT issuer is not configured" + ); + Status::unavailable("sandbox JWT minting is not configured on this gateway") + })?; + + let record = load_sandbox(state, sandbox_id).await?; + let minted = issuer.mint(sandbox_id)?; + let sandbox_name = record + .metadata + .as_ref() + .map_or_else(String::new, |m| m.name.clone()); + info!( + sandbox_id = %sandbox_id, + reason, + "issued gateway sandbox JWT for pod activation" + ); + + Ok(PodActivationMessage { + sandbox_id: sandbox_id.to_string(), + sandbox_name, + token: minted.token, + token_expires_at_ms: minted.expires_at_ms, + startup_metadata: std::collections::HashMap::default(), + }) +} + +pub async fn load_sandbox(state: &Arc, sandbox_id: &str) -> Result { + if sandbox_id.is_empty() { + return Err(Status::invalid_argument("sandbox_id is required")); + } + + state + .store + .get_message::(sandbox_id) + .await + .map_err(|e| Status::internal(format!("fetch sandbox failed: {e}")))? + .ok_or_else(|| Status::not_found("sandbox not found")) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::auth::sandbox_jwt::SandboxJwtIssuer; + use crate::compute::new_test_runtime; + use crate::persistence::Store; + use crate::sandbox_index::SandboxIndex; + use crate::sandbox_watch::SandboxWatchBus; + use crate::supervisor_session::SupervisorSessionRegistry; + use crate::tracing_bus::TracingLogBus; + use openshell_bootstrap::jwt::generate_jwt_key; + use openshell_core::Config; + use openshell_core::proto::datamodel::v1::ObjectMeta; + use openshell_core::proto::{Sandbox, SandboxPhase, SandboxSpec}; + use std::collections::HashMap; + use std::time::Duration; + use tokio_stream::StreamExt; + + struct StaticValidator { + validated: ValidatedWarmPodActivation, + } + + #[async_trait] + impl WarmPodActivationValidator for StaticValidator { + async fn validate( + &self, + _target: &WarmPodActivationTarget, + _pending: &SupervisorBootstrapIdentity, + ) -> Result { + Ok(self.validated.clone()) + } + } + + async fn state_with_issuer() -> Arc { + let mat = generate_jwt_key().expect("jwt key"); + let store = Arc::new( + Store::connect("sqlite::memory:?cache=shared") + .await + .unwrap(), + ); + let compute = new_test_runtime(store.clone()).await; + let mut state = ServerState::new( + Config::new(None).with_database_url("sqlite::memory:?cache=shared"), + store, + compute, + SandboxIndex::new(), + SandboxWatchBus::new(), + TracingLogBus::new(), + Arc::new(SupervisorSessionRegistry::new()), + None, + ); + let issuer = SandboxJwtIssuer::from_pem( + mat.signing_key_pem.as_bytes(), + mat.kid, + "test-gateway", + Duration::from_secs(3600), + ) + .unwrap(); + state.sandbox_jwt_issuer = Some(Arc::new(issuer)); + let state = Arc::new(state); + insert_sandbox(&state, "sandbox-a").await; + state + } + + async fn insert_sandbox(state: &Arc, sandbox_id: &str) { + let mut sandbox = Sandbox { + metadata: Some(ObjectMeta { + id: sandbox_id.to_string(), + name: sandbox_id.to_string(), + created_at_ms: 1_000_000, + labels: HashMap::default(), + annotations: HashMap::default(), + resource_version: 0, + deletion_timestamp_ms: 0, + workspace: "default".to_string(), + }), + spec: Some(SandboxSpec { + policy: None, + ..Default::default() + }), + ..Default::default() + }; + sandbox.set_phase(SandboxPhase::Ready as i32); + state.store.put_message(&sandbox).await.unwrap(); + } + + fn pending_identity() -> SupervisorBootstrapIdentity { + SupervisorBootstrapIdentity { + driver: "kubernetes".to_string(), + instance_name: "warm-pod-a".to_string(), + instance_id: "pod-uid-a".to_string(), + owner_name: "sandbox-owner-a".to_string(), + owner_uid: "owner-uid-a".to_string(), + binding: openshell_core::supervisor_bootstrap::SupervisorBootstrapBinding::WarmPending, + } + } + + fn activation_target(sandbox_id: &str) -> WarmPodActivationTarget { + WarmPodActivationTarget { + driver: "kubernetes".to_string(), + instance_id: "pod-uid-a".to_string(), + sandbox_id: sandbox_id.to_string(), + owner_uid: "owner-uid-a".to_string(), + } + } + + fn validated_activation(sandbox_id: &str) -> ValidatedWarmPodActivation { + ValidatedWarmPodActivation { + driver: "kubernetes".to_string(), + instance_id: "pod-uid-a".to_string(), + instance_name: "warm-pod-a".to_string(), + owner_uid: "owner-uid-a".to_string(), + sandbox_id: sandbox_id.to_string(), + } + } + + #[tokio::test] + async fn pending_warm_pod_receives_activation_token() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = StaticValidator { + validated: validated_activation("sandbox-a"), + }; + + activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect("activate"); + + let received = stream + .next() + .await + .expect("activation message") + .expect("activation OK"); + assert_eq!(received.sandbox_id, "sandbox-a"); + assert_eq!(received.sandbox_name, "sandbox-a"); + assert!(!received.token.is_empty()); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + assert_eq!(state.supervisor_pod_registrations.activated_count(), 1); + } + + #[tokio::test] + async fn activation_for_unknown_instance_id_fails_before_validation() { + let state = state_with_issuer().await; + let validator = StaticValidator { + validated: validated_activation("sandbox-a"), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("unknown instance ID must fail"); + + assert_eq!(err.code(), tonic::Code::NotFound); + } + + #[tokio::test] + async fn missing_target_sandbox_fails_without_token_emission() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = StaticValidator { + validated: validated_activation("sandbox-deleted"), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-deleted")) + .await + .expect_err("missing sandbox must fail"); + + assert_eq!(err.code(), tonic::Code::NotFound); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 1); + assert!( + tokio::time::timeout(Duration::from_millis(10), stream.next()) + .await + .is_err(), + "retryable missing sandbox must leave stream pending" + ); + } + + #[tokio::test] + async fn token_backed_activation_accepts_matching_token() { + let state = state_with_issuer().await; + state + .store + .put_activation_token("sandbox-a", "activation-token-a") + .await + .unwrap(); + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = TokenBackedActivationRequest { + state: state.clone(), + activation_token: "activation-token-a".to_string(), + }; + + activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect("activate"); + + let received = stream + .next() + .await + .expect("activation message") + .expect("activation OK"); + assert_eq!(received.sandbox_id, "sandbox-a"); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + } + + #[tokio::test] + async fn token_backed_activation_missing_mapping_is_retryable() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = TokenBackedActivationRequest { + state: state.clone(), + activation_token: "activation-token-a".to_string(), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("missing mapping must fail"); + + assert_eq!(err.code(), tonic::Code::Unavailable); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 1); + assert!( + tokio::time::timeout(Duration::from_millis(10), stream.next()) + .await + .is_err(), + "retryable token mapping miss must leave stream pending" + ); + } + + #[tokio::test] + async fn token_backed_activation_rejects_token_mismatch_as_fatal() { + let state = state_with_issuer().await; + state + .store + .put_activation_token("sandbox-a", "activation-token-a") + .await + .unwrap(); + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = TokenBackedActivationRequest { + state: state.clone(), + activation_token: "wrong-token".to_string(), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("mismatched token must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + let stream_err = stream + .next() + .await + .expect("failure message") + .expect_err("mismatch should fail stream"); + assert_eq!(stream_err.code(), tonic::Code::PermissionDenied); + } + + #[tokio::test] + async fn token_backed_activation_rejects_empty_token_as_fatal() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = TokenBackedActivationRequest { + state: state.clone(), + activation_token: String::new(), + }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("empty token must fail"); + + assert_eq!(err.code(), tonic::Code::InvalidArgument); + assert_eq!(state.supervisor_pod_registrations.pending_count(), 0); + let stream_err = stream + .next() + .await + .expect("failure message") + .expect_err("empty token should fail stream"); + assert_eq!(stream_err.code(), tonic::Code::InvalidArgument); + } + + #[tokio::test] + async fn duplicate_activation_for_same_instance_id_is_rejected() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let validator = StaticValidator { + validated: validated_activation("sandbox-a"), + }; + + activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect("first activation"); + let _ = stream.next().await.expect("activation"); + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("duplicate activation must fail"); + + assert_eq!(err.code(), tonic::Code::AlreadyExists); + } + + #[tokio::test] + async fn revalidation_metadata_must_match_pending_registration() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register pending"); + let mut validated = validated_activation("sandbox-a"); + validated.owner_uid = "other-owner".to_string(); + let validator = StaticValidator { validated }; + + let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect_err("mismatched revalidation metadata must fail"); + + assert_eq!(err.code(), tonic::Code::PermissionDenied); + let stream_err = stream + .next() + .await + .expect("failure message") + .expect_err("validation mismatch should fail stream"); + assert_eq!(stream_err.code(), tonic::Code::PermissionDenied); + } +} diff --git a/crates/openshell-server/tests/common/mod.rs b/crates/openshell-server/tests/common/mod.rs index 93beeacf15..3352d8d991 100644 --- a/crates/openshell-server/tests/common/mod.rs +++ b/crates/openshell-server/tests/common/mod.rs @@ -21,10 +21,11 @@ use openshell_core::proto::{ GetSandboxConfigResponse, GetSandboxProviderEnvironmentRequest, GetSandboxProviderEnvironmentResponse, GetSandboxRequest, HealthRequest, HealthResponse, IssueSandboxTokenRequest, IssueSandboxTokenResponse, ListProvidersRequest, - ListProvidersResponse, ListSandboxesRequest, ListSandboxesResponse, ProviderResponse, - RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RelayFrame, RevokeSshSessionRequest, - RevokeSshSessionResponse, SandboxResponse, SandboxStreamEvent, ServiceStatus, - SupervisorMessage, TcpForwardFrame, UpdateProviderRequest, WatchSandboxRequest, + ListProvidersResponse, ListSandboxesRequest, ListSandboxesResponse, PodActivationMessage, + ProviderResponse, RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, + RegisterSupervisorPodRequest, RelayFrame, RevokeSshSessionRequest, RevokeSshSessionResponse, + SandboxResponse, SandboxStreamEvent, ServiceStatus, SupervisorMessage, TcpForwardFrame, + UpdateProviderRequest, WatchSandboxRequest, open_shell_client::OpenShellClient, open_shell_server::{OpenShell, OpenShellServer}, }; @@ -448,6 +449,15 @@ impl OpenShell for TestOpenShell { Err(Status::unimplemented("not implemented in test")) } + type RegisterSupervisorPodStream = ReceiverStream>; + + async fn register_supervisor_pod( + &self, + _request: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("not implemented in test")) + } + async fn refresh_sandbox_token( &self, _request: tonic::Request, diff --git a/crates/openshell-server/tests/supervisor_relay_integration.rs b/crates/openshell-server/tests/supervisor_relay_integration.rs index 721baacd88..134e58b840 100644 --- a/crates/openshell-server/tests/supervisor_relay_integration.rs +++ b/crates/openshell-server/tests/supervisor_relay_integration.rs @@ -23,7 +23,8 @@ use hyper_util::{ server::conn::auto::Builder, }; use openshell_core::proto::{ - GatewayMessage, RelayFrame, RelayInit, SupervisorMessage, TcpForwardFrame, + GatewayMessage, PodActivationMessage, RelayFrame, RelayInit, SupervisorMessage, + TcpForwardFrame, open_shell_client::OpenShellClient, open_shell_server::{OpenShell, OpenShellServer}, }; @@ -413,6 +414,13 @@ impl OpenShell for RelayGateway { ) -> Result, Status> { Err(Status::unimplemented("unused")) } + type RegisterSupervisorPodStream = ReceiverStream>; + async fn register_supervisor_pod( + &self, + _: tonic::Request, + ) -> Result, Status> { + Err(Status::unimplemented("unused")) + } async fn refresh_sandbox_token( &self, _: tonic::Request, diff --git a/crates/openshell-supervisor-process/src/run.rs b/crates/openshell-supervisor-process/src/run.rs index 3cb29d1e6b..3c84fbadbc 100644 --- a/crates/openshell-supervisor-process/src/run.rs +++ b/crates/openshell-supervisor-process/src/run.rs @@ -107,6 +107,10 @@ pub async fn run_process( // tasks. By this point the orchestrator has finished privileged startup // helpers (network namespace setup, nftables probes via run_networking), // and the SSH listener and entrypoint child have not been exposed yet. + #[cfg(target_os = "linux")] + if enforcement_mode.uses_privileged_process_setup() { + crate::process::prepare_supervisor_identity_mount_namespace_from_env()?; + } crate::sandbox::apply_supervisor_startup_hardening()?; // Spawn the bypass detection monitor. It tails dmesg for nftables LOG diff --git a/deploy/helm/openshell/README.md b/deploy/helm/openshell/README.md index 8723535d1a..13f7686f34 100644 --- a/deploy/helm/openshell/README.md +++ b/deploy/helm/openshell/README.md @@ -228,6 +228,11 @@ add `ci/values-spire.yaml` to the OpenShell release values files. | server.tls.certSecretName | string | `"openshell-server-tls"` | K8s secret (type kubernetes.io/tls) with tls.crt and tls.key for the server. | | server.tls.clientCaSecretName | string | `"openshell-server-client-ca"` | K8s secret with ca.crt for client certificate verification (mTLS). Set to "" to disable mTLS and run HTTPS-only (use OIDC for auth instead). | | server.tls.clientTlsSecretName | string | `"openshell-client-tls"` | K8s secret mounted into sandbox pods for mTLS to the server. | +| server.warmPooling | object | `{"enabled":true,"profiles":{"enabled":true,"labelSelector":"openshell.ai/warm-pool-profile=true","maxReplicas":100,"namespace":""}}` | Enable transparent Kubernetes warm-pool allocation through Agent Sandbox v1beta1 SandboxClaim resources when a compatible OpenShell-enabled SandboxWarmPool exists in the target namespace. | +| server.warmPooling.profiles | object | `{"enabled":true,"labelSelector":"openshell.ai/warm-pool-profile=true","maxReplicas":100,"namespace":""}` | Reconcile labelled ConfigMap warm-pool profiles into generated SandboxTemplate and SandboxWarmPool resources. Admins can still create warm pools by any other mechanism. | +| server.warmPooling.profiles.labelSelector | string | `"openshell.ai/warm-pool-profile=true"` | Label selector used to select warm-pool profile ConfigMaps. | +| server.warmPooling.profiles.maxReplicas | int | `100` | Maximum replicas accepted from a single warm-pool profile. | +| server.warmPooling.profiles.namespace | string | `""` | Namespace containing warm-pool profile ConfigMaps. Empty defaults to the Helm release namespace. | | server.workspaceDefaultStorageSize | string | `""` | Default storage size for the workspace PVC in sandbox pods. Uses Kubernetes quantity syntax (e.g. "2Gi", "10Gi", "500Mi"). Empty = built-in default (2Gi). | | service.healthPort | int | `8081` | Gateway health service port. | | service.metricsPort | int | `9090` | Gateway metrics service port. | diff --git a/deploy/helm/openshell/templates/clusterrole.yaml b/deploy/helm/openshell/templates/clusterrole.yaml index 073c8835ec..eea70f26ad 100644 --- a/deploy/helm/openshell/templates/clusterrole.yaml +++ b/deploy/helm/openshell/templates/clusterrole.yaml @@ -31,3 +31,31 @@ rules: - namespaces verbs: - get + # Discover OpenShell-enabled Agent Sandbox warm pools across namespaces. + # SandboxClaim create/delete is namespaced, but the ClusterRoleBinding lets + # workspace-to-namespace mappings use warm pools outside the release namespace. + - apiGroups: + - extensions.agents.x-k8s.io + resources: + - sandboxclaims + verbs: + - create + - delete + - get + - list + - watch + - apiGroups: + - extensions.agents.x-k8s.io + resources: + - sandboxtemplates + - sandboxwarmpools + verbs: + - get + - list + - watch + {{- if .Values.server.warmPooling.profiles.enabled }} + - create + - patch + - update + - delete + {{- end }} diff --git a/deploy/helm/openshell/templates/gateway-config.yaml b/deploy/helm/openshell/templates/gateway-config.yaml index 36a579250b..486dc7bda6 100644 --- a/deploy/helm/openshell/templates/gateway-config.yaml +++ b/deploy/helm/openshell/templates/gateway-config.yaml @@ -145,6 +145,15 @@ data: supervisor_image_pull_policy = {{ .Values.supervisor.image.pullPolicy | quote }} {{- end }} + [openshell.drivers.kubernetes.warm_pooling] + enabled = {{ .Values.server.warmPooling.enabled }} + + [openshell.drivers.kubernetes.warm_pooling.profiles] + enabled = {{ .Values.server.warmPooling.profiles.enabled }} + namespace = {{ default .Release.Namespace .Values.server.warmPooling.profiles.namespace | quote }} + label_selector = {{ .Values.server.warmPooling.profiles.labelSelector | quote }} + max_replicas = {{ .Values.server.warmPooling.profiles.maxReplicas | default 100 }} + [openshell.drivers.kubernetes.sidecar] proxy_uid = {{ .Values.supervisor.sidecar.proxyUid | default 1337 }} process_binary_aware_network_policy = {{ .Values.supervisor.sidecar.processBinaryAwareNetworkPolicy }} diff --git a/deploy/helm/openshell/templates/warm-pool-profile-role.yaml b/deploy/helm/openshell/templates/warm-pool-profile-role.yaml new file mode 100644 index 0000000000..a01b53e5f0 --- /dev/null +++ b/deploy/helm/openshell/templates/warm-pool-profile-role.yaml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +{{- if .Values.server.warmPooling.profiles.enabled }} +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: {{ include "openshell.fullname" . }}-warm-pool-profiles + namespace: {{ default .Release.Namespace .Values.server.warmPooling.profiles.namespace }} + labels: + {{- include "openshell.labels" . | nindent 4 }} +rules: + - apiGroups: + - "" + resources: + - configmaps + verbs: + - get + - list + - watch + - patch + - apiGroups: + - "" + resources: + - events + verbs: + - create +{{- end }} diff --git a/deploy/helm/openshell/templates/warm-pool-profile-rolebinding.yaml b/deploy/helm/openshell/templates/warm-pool-profile-rolebinding.yaml new file mode 100644 index 0000000000..a3c82c98be --- /dev/null +++ b/deploy/helm/openshell/templates/warm-pool-profile-rolebinding.yaml @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +{{- if .Values.server.warmPooling.profiles.enabled }} +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: {{ include "openshell.fullname" . }}-warm-pool-profiles + namespace: {{ default .Release.Namespace .Values.server.warmPooling.profiles.namespace }} + labels: + {{- include "openshell.labels" . | nindent 4 }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: {{ include "openshell.fullname" . }}-warm-pool-profiles +subjects: + - kind: ServiceAccount + name: {{ include "openshell.serviceAccountName" . }} + namespace: {{ .Release.Namespace }} +{{- end }} diff --git a/deploy/helm/openshell/tests/gateway_config_test.yaml b/deploy/helm/openshell/tests/gateway_config_test.yaml index aee396c38f..3c4d0e7f30 100644 --- a/deploy/helm/openshell/tests/gateway_config_test.yaml +++ b/deploy/helm/openshell/tests/gateway_config_test.yaml @@ -159,6 +159,36 @@ tests: path: data["gateway.toml"] pattern: '(?ms)\[openshell\.drivers\.kubernetes\.sidecar\].*?process_binary_aware_network_policy\s*=\s*false' + - it: renders warm pooling config under [openshell.drivers.kubernetes.warm_pooling] + template: templates/gateway-config.yaml + set: + server.warmPooling.enabled: false + asserts: + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\].*?enabled\s*=\s*false' + + - it: renders warm pool profile config under [openshell.drivers.kubernetes.warm_pooling.profiles] + template: templates/gateway-config.yaml + set: + server.warmPooling.profiles.enabled: true + server.warmPooling.profiles.namespace: warm-profile-ns + server.warmPooling.profiles.labelSelector: openshell.ai/warm-pool-profile=true,tier=gpu + server.warmPooling.profiles.maxReplicas: 7 + asserts: + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\.profiles\].*?enabled\s*=\s*true' + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\.profiles\].*?namespace\s*=\s*"warm-profile-ns"' + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\.profiles\].*?label_selector\s*=\s*"openshell.ai/warm-pool-profile=true,tier=gpu"' + - matchRegex: + path: data["gateway.toml"] + pattern: '(?ms)\[openshell\.drivers\.kubernetes\.warm_pooling\.profiles\].*?max_replicas\s*=\s*7' + - it: renders sandbox image pull secrets under [openshell.drivers.kubernetes] template: templates/gateway-config.yaml set: diff --git a/deploy/helm/openshell/tests/warm_pool_profiles_rbac_test.yaml b/deploy/helm/openshell/tests/warm_pool_profiles_rbac_test.yaml new file mode 100644 index 0000000000..6815a3e664 --- /dev/null +++ b/deploy/helm/openshell/tests/warm_pool_profiles_rbac_test.yaml @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +suite: warm pool profile RBAC +templates: + - templates/clusterrole.yaml + - templates/warm-pool-profile-role.yaml + - templates/warm-pool-profile-rolebinding.yaml +release: + name: openshell + namespace: gateway-ns + +tests: + - it: does not render profile namespace RBAC when profile reconciliation is explicitly disabled + template: templates/warm-pool-profile-role.yaml + set: + server.warmPooling.profiles.enabled: false + asserts: + - hasDocuments: + count: 0 + + - it: renders profile ConfigMap RBAC in the release namespace by default + template: templates/warm-pool-profile-role.yaml + asserts: + - equal: + path: metadata.namespace + value: gateway-ns + - contains: + path: rules[0].resources + content: configmaps + - contains: + path: rules[0].verbs + content: patch + + - it: renders profile ConfigMap RBAC in the configured namespace + template: templates/warm-pool-profile-rolebinding.yaml + set: + server.warmPooling.profiles.namespace: profiles-ns + asserts: + - equal: + path: metadata.namespace + value: profiles-ns + - equal: + path: subjects[0].namespace + value: gateway-ns + + - it: adds generated warm-pool write verbs to the ClusterRole when enabled + template: templates/clusterrole.yaml + asserts: + - contains: + path: rules[4].verbs + content: create + - contains: + path: rules[4].verbs + content: patch + - contains: + path: rules[4].verbs + content: delete diff --git a/deploy/helm/openshell/values.yaml b/deploy/helm/openshell/values.yaml index e89a234912..ceaf7e14e8 100644 --- a/deploy/helm/openshell/values.yaml +++ b/deploy/helm/openshell/values.yaml @@ -193,6 +193,23 @@ server: # Set to a RuntimeClass name (e.g. "kata-containers", "nvidia") to apply it # to all sandboxes that don't explicitly override it. defaultRuntimeClassName: "" + # -- Enable transparent Kubernetes warm-pool allocation through Agent Sandbox + # v1beta1 SandboxClaim resources when a compatible OpenShell-enabled + # SandboxWarmPool exists in the target namespace. + warmPooling: + enabled: true + # -- Reconcile labelled ConfigMap warm-pool profiles into generated + # SandboxTemplate and SandboxWarmPool resources. Admins can still create + # warm pools by any other mechanism. + profiles: + enabled: true + # -- Namespace containing warm-pool profile ConfigMaps. Empty defaults to + # the Helm release namespace. + namespace: "" + # -- Label selector used to select warm-pool profile ConfigMaps. + labelSelector: "openshell.ai/warm-pool-profile=true" + # -- Maximum replicas accepted from a single warm-pool profile. + maxReplicas: 100 # -- gRPC endpoint sandboxes call back into the gateway. Leave empty to derive # it from the chart fullname, release namespace, service port, and # disableTls flag, for example https://openshell.openshell.svc.cluster.local:8080. diff --git a/docs/kubernetes/setup.mdx b/docs/kubernetes/setup.mdx index c2fca827f1..a0bc78141a 100644 --- a/docs/kubernetes/setup.mdx +++ b/docs/kubernetes/setup.mdx @@ -216,12 +216,16 @@ The namespaced Role covers sandbox lifecycle and identity: | `""` | `events` | get, list, watch | | `""` | `pods` | get | -The ClusterRole grants node inspection and token validation: +The ClusterRole grants node inspection, token validation, namespace reads, and +warm-pool extension access: | API Group | Resource | Verbs | |---|---|---| | `authentication.k8s.io` | `tokenreviews` | create | | `""` | `nodes` | get, list, watch | +| `""` | `namespaces` | get | +| `extensions.agents.x-k8s.io` | `sandboxclaims` | create, delete, get, list, watch | +| `extensions.agents.x-k8s.io` | `sandboxtemplates`, `sandboxwarmpools` | get, list, watch | To use an existing ServiceAccount instead of creating one, set `serviceAccount.create=false` and supply its name: diff --git a/docs/reference/gateway-config.mdx b/docs/reference/gateway-config.mdx index 56f5d15348..fc90a0eb3e 100644 --- a/docs/reference/gateway-config.mdx +++ b/docs/reference/gateway-config.mdx @@ -294,6 +294,18 @@ provider_spiffe_workload_api_socket_path = "/spiffe-workload-api/spire-agent.soc # sandbox_uid = 1500 # sandbox_gid = 1500 +[openshell.drivers.kubernetes.warm_pooling] +# When enabled, compatible creates may use v1beta1 SandboxClaim resources +# backed by the driver's cache of OpenShell-enabled SandboxWarmPool resources. +enabled = true + +[openshell.drivers.kubernetes.warm_pooling.profiles] +# Optional ConfigMap-backed reconciler for admin-declared warm pools. +enabled = true +namespace = "openshell" +label_selector = "openshell.ai/warm-pool-profile=true" +max_replicas = 100 + [openshell.drivers.kubernetes.sidecar] # UID used by relaxed long-running network sidecars. Strict process/binary-aware # sidecars run as UID 0 so Kubernetes grants the required /proc inspection diff --git a/docs/reference/sandbox-compute-drivers.mdx b/docs/reference/sandbox-compute-drivers.mdx index 280849b5f6..c44a536f2c 100644 --- a/docs/reference/sandbox-compute-drivers.mdx +++ b/docs/reference/sandbox-compute-drivers.mdx @@ -340,6 +340,41 @@ process/binary identity through `/proc/`. The Kubernetes driver creates namespaced `agents.x-k8s.io` `Sandbox` resources from the Kubernetes SIG Apps [agent-sandbox](https://github.com/kubernetes-sigs/agent-sandbox) project. It detects the served Sandbox API at runtime, caches the selected API version for the gateway process, and uses `v1beta1` when available before falling back to `v1alpha1`, so supported Agent Sandbox installations work without version-specific operator configuration. The Agent Sandbox controller turns those resources into sandbox pods and related storage. +When Kubernetes warm pooling is enabled, the driver can satisfy compatible +create requests by creating `extensions.agents.x-k8s.io/v1beta1` +`SandboxClaim` resources instead of direct `Sandbox` resources. The driver +discovers `SandboxWarmPool` resources labelled `openshell.ai/enabled=true` +across namespaces it can observe, resolves each pool's `SandboxTemplate`, and +matches only against pools in the Kubernetes namespace selected for the +OpenShell workspace. Set +`openshell.drivers.kubernetes.warm_pooling.enabled = false` to force direct +`Sandbox` creation. + +When `openshell.drivers.kubernetes.warm_pooling.profiles.enabled = true`, the +driver reconciles labelled ConfigMaps into generated `SandboxTemplate` and +`SandboxWarmPool` resources. This is enabled by default with the Helm chart. +Each profile ConfigMap must contain +`data["warm-pool.toml"]` with a narrow CLI-shaped schema: + +```toml +version = 1 +workspace = "default" +replicas = 3 +image = "ghcr.io/acme/agent-python:latest" +runtime_class_name = "nvidia" + +[environment] +FOO = "bar" + +[resources] +cpu = "4" +memory = "8Gi" +gpu_count = 1 +``` + +The generated warm pools carry `openshell.ai/enabled=true`, so the existing +matching path discovers them like any externally managed warm pool. + If Agent Sandbox is upgraded in place, restart the OpenShell gateway after the controller and CRD rollout completes so the gateway can detect the served API versions again. diff --git a/e2e/rust/Cargo.toml b/e2e/rust/Cargo.toml index e4d99dc45c..d5ec1364c1 100644 --- a/e2e/rust/Cargo.toml +++ b/e2e/rust/Cargo.toml @@ -87,6 +87,11 @@ name = "host_gateway_alias" path = "tests/host_gateway_alias.rs" required-features = ["e2e-host-gateway"] +[[test]] +name = "kubernetes_warm_pool" +path = "tests/kubernetes_warm_pool.rs" +required-features = ["e2e-kubernetes"] + [[test]] name = "forward_proxy_l7_bypass" path = "tests/forward_proxy_l7_bypass.rs" diff --git a/e2e/rust/src/harness/kubernetes.rs b/e2e/rust/src/harness/kubernetes.rs new file mode 100644 index 0000000000..5be3a6dbae --- /dev/null +++ b/e2e/rust/src/harness/kubernetes.rs @@ -0,0 +1,219 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Lightweight Kubernetes helpers for e2e tests. +//! +//! These helpers intentionally shell out to `kubectl` so the e2e crate does not +//! need a Kubernetes client dependency. They inherit `KUBECONFIG` and, when +//! present, `OPENSHELL_E2E_KUBE_CONTEXT` from the wrapper. + +use std::process::Stdio; +use std::time::{Duration, Instant}; + +use serde_json::Value; +use tokio::process::Command; +use tokio::time::sleep; + +pub async fn kubectl(args: &[&str]) -> Result { + let mut cmd = Command::new("kubectl"); + if let Ok(context) = std::env::var("OPENSHELL_E2E_KUBE_CONTEXT") + && !context.trim().is_empty() + { + cmd.arg("--context").arg(context); + } + cmd.args(args).stdout(Stdio::piped()).stderr(Stdio::piped()); + + let output = cmd + .output() + .await + .map_err(|err| format!("failed to run kubectl {args:?}: {err}"))?; + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + let combined = format!("{stdout}{stderr}"); + + if !output.status.success() { + return Err(format!( + "kubectl {args:?} failed (exit {:?}):\n{combined}", + output.status.code() + )); + } + + Ok(stdout) +} + +pub async fn kubectl_json(args: &[&str]) -> Result { + let output = kubectl(args).await?; + serde_json::from_str(&output) + .map_err(|err| format!("kubectl {args:?} did not return valid JSON: {err}\n{output}")) +} + +pub async fn has_crd(plural: &str, group: &str) -> bool { + let name = format!("{plural}.{group}"); + kubectl(&["get", "crd", &name]).await.is_ok() +} + +pub async fn wait_for_jsonpath( + namespace: &str, + kind: &str, + name: &str, + jsonpath: &str, + expected: &str, + timeout: Duration, +) -> Result<(), String> { + let deadline = Instant::now() + timeout; + let mut last_output: Option; + let output_arg = format!("jsonpath={jsonpath}"); + + loop { + match kubectl(&["get", kind, name, "-n", namespace, "-o", &output_arg]).await { + Ok(output) => { + let trimmed = output.trim().to_string(); + if trimmed == expected { + return Ok(()); + } + last_output = Some(trimmed); + } + Err(err) => last_output = Some(err), + } + + if Instant::now() >= deadline { + let last_output = last_output.unwrap_or_else(|| "".to_string()); + return Err(format!( + "timed out after {}s waiting for {kind}/{name} jsonpath {jsonpath} to equal {expected:?}. Last output:\n{last_output}", + timeout.as_secs() + )); + } + sleep(Duration::from_secs(2)).await; + } +} + +pub async fn wait_for_resource_by_label( + namespace: &str, + kind: &str, + selector: &str, + timeout: Duration, +) -> Result { + let deadline = Instant::now() + timeout; + let mut last_output: Option; + + loop { + match kubectl_json(&["get", kind, "-n", namespace, "-l", selector, "-o", "json"]).await { + Ok(value) if !items(&value).is_empty() => return Ok(value), + Ok(value) => last_output = Some(value.to_string()), + Err(err) => last_output = Some(err), + } + + if Instant::now() >= deadline { + let last_output = last_output.unwrap_or_else(|| "".to_string()); + return Err(format!( + "timed out after {}s waiting for {kind} in namespace {namespace} with selector {selector}. Last output:\n{last_output}", + timeout.as_secs() + )); + } + sleep(Duration::from_secs(2)).await; + } +} + +pub async fn wait_for_resource_absent_by_label( + namespace: &str, + kind: &str, + selector: &str, + timeout: Duration, +) -> Result<(), String> { + let deadline = Instant::now() + timeout; + let mut last_output: Option; + + loop { + match kubectl_json(&["get", kind, "-n", namespace, "-l", selector, "-o", "json"]).await { + Ok(value) if items(&value).is_empty() => return Ok(()), + Ok(value) => last_output = Some(value.to_string()), + Err(err) if err.contains("the server doesn't have a resource type") => return Ok(()), + Err(err) => last_output = Some(err), + } + + if Instant::now() >= deadline { + let last_output = last_output.unwrap_or_else(|| "".to_string()); + return Err(format!( + "timed out after {}s waiting for {kind} in namespace {namespace} with selector {selector} to be absent. Last output:\n{last_output}", + timeout.as_secs() + )); + } + sleep(Duration::from_secs(2)).await; + } +} + +pub async fn delete_resource(namespace: &str, kind: &str, name: &str) -> Result { + kubectl(&[ + "delete", + kind, + name, + "-n", + namespace, + "--ignore-not-found=true", + ]) + .await +} + +pub async fn dump_namespace_diagnostics(namespace: &str) { + eprintln!("=== Kubernetes diagnostics for namespace {namespace} ==="); + for args in [ + vec![ + "get", + "sandboxclaims.extensions.agents.x-k8s.io", + "-n", + namespace, + "-o", + "wide", + ], + vec![ + "get", + "sandboxwarmpools.extensions.agents.x-k8s.io", + "-n", + namespace, + "-o", + "wide", + ], + vec![ + "get", + "sandboxtemplates.extensions.agents.x-k8s.io", + "-n", + namespace, + "-o", + "wide", + ], + vec![ + "get", + "sandboxes.agents.x-k8s.io", + "-n", + namespace, + "-o", + "wide", + ], + vec!["get", "pods", "-n", namespace, "-o", "wide"], + vec!["get", "events", "-n", namespace, "--sort-by=.lastTimestamp"], + vec![ + "logs", + "-n", + namespace, + "-l", + "app.kubernetes.io/instance=openshell", + "--tail=200", + "--all-containers", + "--prefix", + ], + ] { + eprintln!("--- kubectl {args:?} ---"); + match kubectl(&args).await { + Ok(output) => eprintln!("{output}"), + Err(err) => eprintln!("{err}"), + } + } + eprintln!("=== end Kubernetes diagnostics ==="); +} + +pub fn items(value: &Value) -> &[Value] { + value + .get("items") + .and_then(Value::as_array) + .map_or(&[], Vec::as_slice) +} diff --git a/e2e/rust/src/harness/mod.rs b/e2e/rust/src/harness/mod.rs index f2dfd5ec9c..272e1e303b 100644 --- a/e2e/rust/src/harness/mod.rs +++ b/e2e/rust/src/harness/mod.rs @@ -7,6 +7,7 @@ pub mod binary; pub mod cli; pub mod container; pub mod gateway; +pub mod kubernetes; pub mod output; pub mod port; pub mod sandbox; diff --git a/e2e/rust/tests/kubernetes_warm_pool.rs b/e2e/rust/tests/kubernetes_warm_pool.rs new file mode 100644 index 0000000000..290bedb918 --- /dev/null +++ b/e2e/rust/tests/kubernetes_warm_pool.rs @@ -0,0 +1,671 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![cfg(feature = "e2e-kubernetes")] + +use std::io::Write; +use std::time::Instant; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use openshell_e2e::harness::cli::{run_cli, sandbox_names, wait_for_sandbox_exec_contains}; +use openshell_e2e::harness::kubernetes::{ + delete_resource, dump_namespace_diagnostics, has_crd, items, kubectl, kubectl_json, + wait_for_jsonpath, wait_for_resource_absent_by_label, wait_for_resource_by_label, +}; +use serde_json::Value; +use tempfile::NamedTempFile; + +const EXTENSIONS_GROUP: &str = "extensions.agents.x-k8s.io"; +const AGENTS_GROUP: &str = "agents.x-k8s.io"; +const CLAIM_KIND: &str = "sandboxclaims.extensions.agents.x-k8s.io"; +const WARM_POOL_KIND: &str = "sandboxwarmpools.extensions.agents.x-k8s.io"; +const TEMPLATE_KIND: &str = "sandboxtemplates.extensions.agents.x-k8s.io"; +const SANDBOX_KIND: &str = "sandboxes.agents.x-k8s.io"; + +const LABEL_ENABLED: &str = "openshell.ai/enabled"; +const LABEL_MANAGED_BY: &str = "openshell.ai/managed-by"; +const LABEL_PROFILE_UID: &str = "openshell.ai/warm-pool-profile-uid"; +const LABEL_ALLOCATION: &str = "openshell.ai/allocation"; +const LABEL_SANDBOX_NAME: &str = "openshell.ai/sandbox-name"; +const LABEL_SANDBOX_WORKSPACE: &str = "openshell.ai/sandbox-workspace"; + +const ANNOTATION_STATUS: &str = "openshell.ai/warm-pool-status"; +const ANNOTATION_TEMPLATE: &str = "openshell.ai/warm-pool-template"; +const ANNOTATION_FINGERPRINT: &str = "openshell.ai/warm-pool-fingerprint"; + +#[derive(Clone, Copy)] +enum ProfileCase { + Default, + Env, + Cpu, +} + +struct GeneratedProfile { + source_name: String, + selector: String, + template_name: String, + warm_pool_name: String, +} + +struct TestContext { + namespace: String, + profiles: Vec, + sandboxes: Vec, +} + +impl TestContext { + fn new(namespace: String) -> Self { + Self { + namespace, + profiles: Vec::new(), + sandboxes: Vec::new(), + } + } + + async fn cleanup(&mut self) { + for sandbox in self.sandboxes.iter().rev() { + let _ = run_cli(&["sandbox", "delete", sandbox]).await; + } + for profile in self.profiles.iter().rev() { + let _ = delete_resource(&self.namespace, "configmap", profile).await; + } + } +} + +#[tokio::test] +async fn kubernetes_warm_pool_profiles_claim_and_fallback() { + if !env_flag("OPENSHELL_E2E_KUBE_WARM_POOL") { + eprintln!( + "Skipping Kubernetes warm-pool e2e test: set OPENSHELL_E2E_KUBE_WARM_POOL=1 to enable" + ); + return; + } + + let namespace = + std::env::var("OPENSHELL_E2E_SANDBOX_NAMESPACE").unwrap_or_else(|_| "openshell".into()); + let strict = env_flag("OPENSHELL_E2E_KUBE_WARM_POOL_STRICT"); + if let Err(err) = ensure_required_crds(strict).await { + if strict { + panic!("{err}"); + } + eprintln!("Skipping Kubernetes warm-pool e2e test: {err}"); + return; + } + + let mut ctx = TestContext::new(namespace); + let result = run_warm_pool_e2e(&mut ctx).await; + if let Err(err) = result { + dump_namespace_diagnostics(&ctx.namespace).await; + ctx.cleanup().await; + panic!("{err}"); + } + ctx.cleanup().await; +} + +async fn run_warm_pool_e2e(ctx: &mut TestContext) -> Result<(), String> { + let profile_suffix = unique_profile_suffix(); + let sandbox_suffix = unique_sandbox_suffix(); + + let default_profile = apply_and_assert_profile( + ctx, + &format!("openshell-wp-e2e-default-{profile_suffix}"), + ProfileCase::Default, + ) + .await?; + let env_profile = apply_and_assert_profile( + ctx, + &format!("openshell-wp-e2e-env-{profile_suffix}"), + ProfileCase::Env, + ) + .await?; + let cpu_profile = apply_and_assert_profile( + ctx, + &format!("openshell-wp-e2e-cpu-{profile_suffix}"), + ProfileCase::Cpu, + ) + .await?; + + create_and_assert_claimed( + ctx, + &format!("wpd-{sandbox_suffix}"), + &[], + "warm-pool-ok", + &default_profile, + ) + .await?; + create_and_assert_claimed( + ctx, + &format!("wpe-{sandbox_suffix}"), + &["--env", "FOO=bar"], + "env-ok", + &env_profile, + ) + .await?; + create_and_assert_claimed( + ctx, + &format!("wpc-{sandbox_suffix}"), + &["--cpu", "0.2"], + "cpu-ok", + &cpu_profile, + ) + .await?; + + create_and_assert_direct_fallback(ctx, &format!("wpf-{sandbox_suffix}")).await?; + delete_profile_and_assert_gc(ctx, &env_profile, &[&default_profile, &cpu_profile]).await?; + + Ok(()) +} + +async fn ensure_required_crds(strict: bool) -> Result<(), String> { + let required = [ + ("sandboxes", AGENTS_GROUP), + ("sandboxclaims", EXTENSIONS_GROUP), + ("sandboxtemplates", EXTENSIONS_GROUP), + ("sandboxwarmpools", EXTENSIONS_GROUP), + ]; + let mut missing = Vec::new(); + for (plural, group) in required { + if !has_crd(plural, group).await { + missing.push(format!("{plural}.{group}")); + } + } + if missing.is_empty() { + Ok(()) + } else if strict { + Err(format!( + "required Agent Sandbox CRDs are missing in strict mode: {}", + missing.join(", ") + )) + } else { + Err(format!( + "required Agent Sandbox CRDs are missing: {}", + missing.join(", ") + )) + } +} + +async fn apply_and_assert_profile( + ctx: &mut TestContext, + name: &str, + case: ProfileCase, +) -> Result { + ctx.profiles.push(name.to_string()); + let manifest = profile_configmap_yaml(&ctx.namespace, name, case); + let mut file = + NamedTempFile::new().map_err(|err| format!("create profile manifest tempfile: {err}"))?; + file.write_all(manifest.as_bytes()) + .map_err(|err| format!("write profile manifest tempfile: {err}"))?; + let path = file + .path() + .to_str() + .ok_or_else(|| "profile manifest tempfile path is not UTF-8".to_string())?; + kubectl(&["apply", "-f", path]).await?; + + wait_for_jsonpath( + &ctx.namespace, + "configmap", + name, + "{.metadata.annotations.openshell\\.ai/warm-pool-status}", + "Ready", + Duration::from_secs(120), + ) + .await?; + + let source = + kubectl_json(&["get", "configmap", name, "-n", &ctx.namespace, "-o", "json"]).await?; + let source_uid = source["metadata"]["uid"] + .as_str() + .ok_or_else(|| format!("ConfigMap/{name} is missing metadata.uid"))? + .to_string(); + let annotations = source["metadata"]["annotations"] + .as_object() + .ok_or_else(|| format!("ConfigMap/{name} is missing status annotations"))?; + let template_annotation = annotations + .get(ANNOTATION_TEMPLATE) + .and_then(Value::as_str) + .ok_or_else(|| format!("ConfigMap/{name} is missing {ANNOTATION_TEMPLATE} annotation"))?; + let fingerprint = annotations + .get(ANNOTATION_FINGERPRINT) + .and_then(Value::as_str) + .ok_or_else(|| format!("ConfigMap/{name} is missing {ANNOTATION_FINGERPRINT} annotation"))? + .to_string(); + if fingerprint.is_empty() { + return Err(format!( + "ConfigMap/{name} has empty {ANNOTATION_FINGERPRINT} annotation" + )); + } + if annotations.get(ANNOTATION_STATUS).and_then(Value::as_str) != Some("Ready") { + return Err(format!( + "ConfigMap/{name} did not reconcile Ready: {source}" + )); + } + + let selector = + format!("{LABEL_MANAGED_BY}=openshell-kubernetes-driver,{LABEL_PROFILE_UID}={source_uid}"); + let template = single_labelled_resource(&ctx.namespace, TEMPLATE_KIND, &selector).await?; + let warm_pool = single_labelled_resource(&ctx.namespace, WARM_POOL_KIND, &selector).await?; + + let template_name = object_name(&template)?; + let warm_pool_name = object_name(&warm_pool)?; + if template_name != template_annotation || warm_pool_name != template_annotation { + return Err(format!( + "generated resource names should match {ANNOTATION_TEMPLATE}={template_annotation}; template={template_name}, warm_pool={warm_pool_name}" + )); + } + + assert_generated_labels(&template, name)?; + assert_generated_labels(&warm_pool, name)?; + if warm_pool["spec"]["sandboxTemplateRef"]["name"].as_str() != Some(&template_name) { + return Err(format!( + "SandboxWarmPool/{warm_pool_name} does not reference SandboxTemplate/{template_name}: {warm_pool}" + )); + } + assert_template_invariants(&template, case)?; + wait_for_warm_pool_ready(&ctx.namespace, &warm_pool_name, Duration::from_secs(120)).await?; + + Ok(GeneratedProfile { + source_name: name.to_string(), + selector, + template_name, + warm_pool_name, + }) +} + +async fn single_labelled_resource( + namespace: &str, + kind: &str, + selector: &str, +) -> Result { + let list = + wait_for_resource_by_label(namespace, kind, selector, Duration::from_secs(120)).await?; + let list_items = items(&list); + if list_items.len() != 1 { + return Err(format!( + "expected exactly one {kind} with selector {selector}, got {}: {list}", + list_items.len() + )); + } + Ok(list_items[0].clone()) +} + +async fn wait_for_warm_pool_ready( + namespace: &str, + name: &str, + timeout: Duration, +) -> Result<(), String> { + let deadline = Instant::now() + timeout; + let mut last_observation: Option; + + loop { + match kubectl_json(&["get", WARM_POOL_KIND, name, "-n", namespace, "-o", "json"]).await { + Ok(value) if warm_pool_ready(&value) => return Ok(()), + Ok(value) => last_observation = Some(value), + Err(err) => last_observation = Some(serde_json::json!({ "error": err })), + } + + if Instant::now() >= deadline { + let last_observation = + last_observation.unwrap_or_else(|| serde_json::json!("")); + return Err(format!( + "SandboxWarmPool/{name} did not become ready within {}s. Last observation:\n{}", + timeout.as_secs(), + serde_json::to_string_pretty(&last_observation) + .unwrap_or_else(|_| last_observation.to_string()) + )); + } + tokio::time::sleep(Duration::from_secs(2)).await; + } +} + +fn warm_pool_ready(value: &Value) -> bool { + let desired = value["spec"]["replicas"].as_u64().unwrap_or(1); + if desired == 0 { + return true; + } + + if status_condition_is_true(value, "Ready") || status_condition_is_true(value, "Available") { + return true; + } + + for path in [ + "/status/readyReplicas", + "/status/availableReplicas", + "/status/currentReadyReplicas", + ] { + if value + .pointer(path) + .and_then(Value::as_u64) + .is_some_and(|ready| ready >= desired) + { + return true; + } + } + + false +} + +fn status_condition_is_true(value: &Value, condition_type: &str) -> bool { + value["status"]["conditions"] + .as_array() + .into_iter() + .flatten() + .any(|condition| { + condition["type"].as_str() == Some(condition_type) + && condition["status"] + .as_str() + .is_some_and(|status| status.eq_ignore_ascii_case("true")) + }) +} + +async fn create_and_assert_claimed( + ctx: &mut TestContext, + sandbox_name: &str, + extra_create_args: &[&str], + marker: &str, + profile: &GeneratedProfile, +) -> Result<(), String> { + ctx.sandboxes.push(sandbox_name.to_string()); + let mut args = vec!["sandbox", "create", "--name", sandbox_name, "--no-tty"]; + args.extend_from_slice(extra_create_args); + args.extend_from_slice(&["--", "echo", marker]); + let (output, code) = run_cli(&args).await; + if code != 0 || !output.contains(marker) { + return Err(format!( + "sandbox create for {sandbox_name} did not succeed through warm pool (exit {code}); expected marker {marker:?}. Output:\n{output}" + )); + } + + let claim = wait_for_claim(ctx, sandbox_name).await?; + if claim["metadata"]["labels"][LABEL_ALLOCATION].as_str() != Some("sandbox-claim") { + return Err(format!( + "SandboxClaim for {sandbox_name} is missing allocation label: {claim}" + )); + } + if claim["metadata"]["labels"][LABEL_SANDBOX_WORKSPACE].as_str() != Some("default") { + return Err(format!( + "SandboxClaim for {sandbox_name} is not in default workspace: {claim}" + )); + } + if claim["spec"]["warmPoolRef"]["name"].as_str() != Some(&profile.warm_pool_name) { + return Err(format!( + "SandboxClaim for {sandbox_name} used wrong warm pool; expected {}, claim: {claim}", + profile.warm_pool_name + )); + } + + let names = sandbox_names().await?; + if !names.iter().any(|name| name == sandbox_name) { + return Err(format!( + "sandbox list --names did not include claimed sandbox {sandbox_name}; names={names:?}" + )); + } + wait_for_sandbox_exec_contains( + sandbox_name, + &["echo", "claimed-ok"], + "claimed-ok", + Duration::from_secs(120), + ) + .await?; + + Ok(()) +} + +async fn wait_for_claim(ctx: &TestContext, sandbox_name: &str) -> Result { + let selector = format!("{LABEL_ALLOCATION}=sandbox-claim,{LABEL_SANDBOX_NAME}={sandbox_name}"); + single_labelled_resource(&ctx.namespace, CLAIM_KIND, &selector).await +} + +async fn create_and_assert_direct_fallback( + ctx: &mut TestContext, + sandbox_name: &str, +) -> Result<(), String> { + ctx.sandboxes.push(sandbox_name.to_string()); + let args = [ + "sandbox", + "create", + "--name", + sandbox_name, + "--no-tty", + "--env", + "FOO=baz", + "--", + "echo", + "fallback-ok", + ]; + let (output, code) = run_cli(&args).await; + if code != 0 || !output.contains("fallback-ok") { + return Err(format!( + "fallback sandbox create failed (exit {code}); output:\n{output}" + )); + } + + let claim_selector = + format!("{LABEL_ALLOCATION}=sandbox-claim,{LABEL_SANDBOX_NAME}={sandbox_name}"); + let claims = kubectl_json(&[ + "get", + CLAIM_KIND, + "-n", + &ctx.namespace, + "-l", + &claim_selector, + "-o", + "json", + ]) + .await?; + if !items(&claims).is_empty() { + return Err(format!( + "fallback sandbox {sandbox_name} unexpectedly allocated through SandboxClaim: {claims}" + )); + } + + let sandbox_selector = + format!("{LABEL_MANAGED_BY}=openshell,{LABEL_SANDBOX_NAME}={sandbox_name}"); + let sandbox = single_labelled_resource(&ctx.namespace, SANDBOX_KIND, &sandbox_selector).await?; + if object_name(&sandbox)? != format!("default--{sandbox_name}") { + return Err(format!( + "fallback sandbox used unexpected direct Sandbox name: {sandbox}" + )); + } + + Ok(()) +} + +async fn delete_profile_and_assert_gc( + ctx: &mut TestContext, + deleted: &GeneratedProfile, + remaining: &[&GeneratedProfile], +) -> Result<(), String> { + delete_resource(&ctx.namespace, "configmap", &deleted.source_name).await?; + wait_for_resource_absent_by_label( + &ctx.namespace, + WARM_POOL_KIND, + &deleted.selector, + Duration::from_secs(120), + ) + .await?; + wait_for_resource_absent_by_label( + &ctx.namespace, + TEMPLATE_KIND, + &deleted.selector, + Duration::from_secs(120), + ) + .await?; + + for profile in remaining { + let template = + single_labelled_resource(&ctx.namespace, TEMPLATE_KIND, &profile.selector).await?; + let warm_pool = + single_labelled_resource(&ctx.namespace, WARM_POOL_KIND, &profile.selector).await?; + if object_name(&template)? != profile.template_name + || object_name(&warm_pool)? != profile.warm_pool_name + { + return Err(format!( + "generated resources for profile {} changed while deleting {}", + profile.source_name, deleted.source_name + )); + } + } + + Ok(()) +} + +fn profile_configmap_yaml(namespace: &str, name: &str, case: ProfileCase) -> String { + let mut profile = String::from( + r#" version = 1 + workspace = "default" + replicas = 1 +"#, + ); + match case { + ProfileCase::Default => {} + ProfileCase::Env => profile.push_str( + r#" + [environment] + FOO = "bar" +"#, + ), + ProfileCase::Cpu => profile.push_str( + r#" + [resources] + cpu = "0.2" +"#, + ), + } + + format!( + r#"apiVersion: v1 +kind: ConfigMap +metadata: + name: {name} + namespace: {namespace} + labels: + openshell.ai/warm-pool-profile: "true" +data: + warm-pool.toml: | +{profile}"# + ) +} + +fn assert_generated_labels(obj: &Value, source_name: &str) -> Result<(), String> { + let labels = obj["metadata"]["labels"] + .as_object() + .ok_or_else(|| format!("generated object is missing metadata.labels: {obj}"))?; + if labels.get(LABEL_ENABLED).and_then(Value::as_str) != Some("true") { + return Err(format!( + "generated object is missing {LABEL_ENABLED}=true: {obj}" + )); + } + if labels.get(LABEL_MANAGED_BY).and_then(Value::as_str) != Some("openshell-kubernetes-driver") { + return Err(format!( + "generated object is missing driver manager label: {obj}" + )); + } + if labels + .get("openshell.ai/warm-pool-profile") + .and_then(Value::as_str) + .is_none_or(str::is_empty) + { + return Err(format!( + "generated object for source {source_name} is missing profile label: {obj}" + )); + } + Ok(()) +} + +fn assert_template_invariants(template: &Value, case: ProfileCase) -> Result<(), String> { + if template["spec"]["podTemplate"]["spec"]["dnsPolicy"].as_str() != Some("ClusterFirst") { + return Err(format!( + "generated SandboxTemplate must set pod dnsPolicy=ClusterFirst: {template}" + )); + } + + let agent = agent_container(template)?; + let image = agent["image"] + .as_str() + .ok_or_else(|| format!("agent container is missing image: {template}"))?; + if image.is_empty() { + return Err(format!( + "agent container image must not be empty: {template}" + )); + } + + let env = agent["env"].as_array().map_or(&[][..], Vec::as_slice); + for reserved in ["OPENSHELL_SANDBOX_ID", "OPENSHELL_SANDBOX"] { + if env + .iter() + .any(|entry| entry["name"].as_str() == Some(reserved)) + { + return Err(format!( + "warm-pool template must not include reserved env {reserved}: {template}" + )); + } + } + + match case { + ProfileCase::Default => {} + ProfileCase::Env => { + let foo = env + .iter() + .find(|entry| entry["name"].as_str() == Some("FOO")) + .and_then(|entry| entry["value"].as_str()); + if foo != Some("bar") { + return Err(format!( + "env warm-pool template did not include FOO=bar: {template}" + )); + } + } + ProfileCase::Cpu => { + let resources = &agent["resources"]; + if resources["limits"]["cpu"].as_str() != Some("0.2") + || resources["requests"]["cpu"].as_str() != Some("0.2") + { + return Err(format!( + "CPU warm-pool template did not render cpu requests/limits=0.2: {template}" + )); + } + } + } + + Ok(()) +} + +fn agent_container(template: &Value) -> Result<&Value, String> { + let containers = template["spec"]["podTemplate"]["spec"]["containers"] + .as_array() + .ok_or_else(|| format!("SandboxTemplate is missing pod containers: {template}"))?; + containers + .iter() + .find(|container| container["name"].as_str() == Some("agent")) + .ok_or_else(|| format!("SandboxTemplate is missing agent container: {template}")) +} + +fn object_name(obj: &Value) -> Result { + obj["metadata"]["name"] + .as_str() + .map(ToOwned::to_owned) + .ok_or_else(|| format!("object is missing metadata.name: {obj}")) +} + +fn unique_profile_suffix() -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + let value = (nanos ^ u128::from(std::process::id())) & 0xffff_ffff; + format!("{value:08x}") +} + +fn unique_sandbox_suffix() -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + let value = (nanos ^ u128::from(std::process::id())) & 0x00ff_ffff; + format!("{value:06x}") +} + +fn env_flag(name: &str) -> bool { + std::env::var(name).is_ok_and(|value| { + value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes") + }) +} diff --git a/e2e/with-kube-gateway.sh b/e2e/with-kube-gateway.sh index 0a114288e8..2e585abaf3 100755 --- a/e2e/with-kube-gateway.sh +++ b/e2e/with-kube-gateway.sh @@ -593,6 +593,10 @@ fi echo "Installing agent-sandbox CRDs and controller (${AGENT_SANDBOX_VERSION})..." _agent_sandbox_base="https://github.com/kubernetes-sigs/agent-sandbox/releases/download/${AGENT_SANDBOX_VERSION}" kctl apply -f "${_agent_sandbox_base}/manifest.yaml" +if [ "${OPENSHELL_E2E_KUBE_WARM_POOL:-0}" = "1" ]; then + echo "Installing agent-sandbox extension CRDs and controllers (${AGENT_SANDBOX_VERSION})..." + kctl apply -f "${_agent_sandbox_base}/extensions.yaml" +fi wait_for_agent_sandbox_crd kctl -n agent-sandbox-system rollout status deployment/agent-sandbox-controller --timeout=300s diff --git a/examples/kubernetes-warm-pool-config/cpu-0-2-configmap.yaml b/examples/kubernetes-warm-pool-config/cpu-0-2-configmap.yaml new file mode 100644 index 0000000000..264d94e8ed --- /dev/null +++ b/examples/kubernetes-warm-pool-config/cpu-0-2-configmap.yaml @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +apiVersion: v1 +kind: ConfigMap +metadata: + name: openshell-warm-pool-cpu-0-2 + namespace: openshell + labels: + openshell.ai/warm-pool-profile: "true" +data: + warm-pool.toml: | + version = 1 + workspace = "openshell" + replicas = 1 + + [resources] + cpu = "0.2" diff --git a/examples/kubernetes-warm-pool-config/default-configmap.yaml b/examples/kubernetes-warm-pool-config/default-configmap.yaml new file mode 100644 index 0000000000..46050ba93f --- /dev/null +++ b/examples/kubernetes-warm-pool-config/default-configmap.yaml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +apiVersion: v1 +kind: ConfigMap +metadata: + name: openshell-warm-pool-default + namespace: openshell + labels: + openshell.ai/warm-pool-profile: "true" +data: + warm-pool.toml: | + version = 1 + workspace = "openshell" + replicas = 1 diff --git a/examples/kubernetes-warm-pool-config/env-foo-configmap.yaml b/examples/kubernetes-warm-pool-config/env-foo-configmap.yaml new file mode 100644 index 0000000000..608f0d5d85 --- /dev/null +++ b/examples/kubernetes-warm-pool-config/env-foo-configmap.yaml @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +apiVersion: v1 +kind: ConfigMap +metadata: + name: openshell-warm-pool-env-foo + namespace: openshell + labels: + openshell.ai/warm-pool-profile: "true" +data: + warm-pool.toml: | + version = 1 + workspace = "openshell" + replicas = 1 + + [environment] + FOO = "bar" diff --git a/proto/compute_driver.proto b/proto/compute_driver.proto index c99b7756b0..55fa888c6f 100644 --- a/proto/compute_driver.proto +++ b/proto/compute_driver.proto @@ -242,7 +242,12 @@ message CreateSandboxRequest { DriverSandbox sandbox = 1; } -message CreateSandboxResponse {} +message CreateSandboxResponse { + // Opaque driver-issued token that must be presented during delayed warm + // activation. Empty for cold creates and drivers that do not use delayed + // activation. + string activation_token = 1 [(openshell.options.v1.secret) = true]; +} message StopSandboxRequest { // Stable sandbox ID stored by the gateway. diff --git a/proto/openshell.proto b/proto/openshell.proto index 1b447a17e3..030fa8be48 100644 --- a/proto/openshell.proto +++ b/proto/openshell.proto @@ -240,6 +240,14 @@ service OpenShell { // and never call this RPC. rpc IssueSandboxToken(IssueSandboxTokenRequest) returns (IssueSandboxTokenResponse); + // Register a Kubernetes supervisor pod and wait for gateway activation. + // + // The initial trial activates already-bound cold pods immediately. Later + // warm-pool stages keep this stream pending until a SandboxClaim adopts the + // registered pod. + rpc RegisterSupervisorPod(RegisterSupervisorPodRequest) + returns (stream PodActivationMessage); + // Renew the calling sandbox's gateway JWT. Older tokens remain valid // until their own expiry; deployments should keep token TTLs short to // bound replay exposure. The supervisor calls this from a background @@ -291,6 +299,28 @@ message IssueSandboxTokenResponse { int64 expires_at_ms = 2; } +// RegisterSupervisorPod request. Empty body; identity is established by the +// authentication credentials carried in the request headers (a projected +// Kubernetes ServiceAccount JWT in the K8s driver path). +message RegisterSupervisorPodRequest {} + +// Activation sent by the gateway once a registered supervisor pod is bound to +// an OpenShell sandbox identity. +message PodActivationMessage { + // OpenShell sandbox UUID the supervisor should use for ConnectSupervisor. + string sandbox_id = 1; + // Human-readable sandbox name, if known. + string sandbox_name = 2; + // Gateway-minted JWT bound to the sandbox UUID. + string token = 3; + // Absolute expiry of the issued token, milliseconds since the epoch. 0 means + // the token is non-expiring. + int64 token_expires_at_ms = 4; + // Reserved for future activation-time metadata. The phase-1 cold path leaves + // this empty. + map startup_metadata = 5; +} + // RefreshSandboxToken request. Empty body; the calling principal must // already be a sandbox principal (i.e. the request carries a still-valid // gateway-minted JWT in its Authorization header). diff --git a/tasks/python.toml b/tasks/python.toml index e04bf27bb4..22e4e1f3c3 100644 --- a/tasks/python.toml +++ b/tasks/python.toml @@ -221,7 +221,7 @@ description = "Generate Python protobuf stubs from .proto files" run = """ #!/usr/bin/env bash set -euo pipefail -uv run --frozen python -m grpc_tools.protoc \ +uv run --frozen --group dev python -m grpc_tools.protoc \ -Iproto \ --python_out=python/openshell/_proto \ --pyi_out=python/openshell/_proto \ @@ -232,7 +232,7 @@ uv run --frozen python -m grpc_tools.protoc \ proto/options.proto \ proto/sandbox.proto # Fix absolute imports in generated stubs to use package-relative imports -uv run --frozen python - <<'PY' +uv run --frozen --group dev python - <<'PY' from pathlib import Path import re