From 84be163b4d4d8ed13e24f563dff93a75f89eec47 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Fri, 17 Jul 2026 16:03:06 +0100 Subject: [PATCH 01/15] fix(python): make proto generation install grpcio tools Signed-off-by: Gordon Sim --- tasks/python.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 From 98e00bf8d107a6ec60b109f3e3aec13a3c0c0b34 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Fri, 17 Jul 2026 16:04:40 +0100 Subject: [PATCH 02/15] feat(k8s): bootstrap supervisors through pod registration stream Signed-off-by: Gordon Sim --- architecture/compute-runtimes.md | 9 + architecture/gateway.md | 28 +-- .../tests/ensure_providers_integration.rs | 11 ++ .../openshell-cli/tests/mtls_integration.rs | 11 ++ .../tests/provider_commands_integration.rs | 11 ++ .../sandbox_create_lifecycle_integration.rs | 11 ++ .../sandbox_name_fallback_integration.rs | 11 ++ crates/openshell-core/src/grpc_client.rs | 36 ++-- crates/openshell-core/src/sandbox_env.rs | 7 +- crates/openshell-driver-kubernetes/README.md | 5 +- .../openshell-driver-kubernetes/src/config.rs | 6 +- .../openshell-driver-kubernetes/src/driver.rs | 13 +- .../openshell-driver-kubernetes/src/main.rs | 4 +- crates/openshell-sdk/tests/client_mock.rs | 10 + .../src/auth/authenticator.rs | 2 +- crates/openshell-server/src/auth/k8s_sa.rs | 44 +++-- crates/openshell-server/src/auth/principal.rs | 4 +- .../openshell-server/src/auth/sandbox_jwt.rs | 2 +- .../src/auth/sandbox_methods.rs | 3 + crates/openshell-server/src/config_file.rs | 2 +- crates/openshell-server/src/grpc/auth_rpc.rs | 181 +++++++++++++----- crates/openshell-server/src/grpc/mod.rs | 33 +++- crates/openshell-server/src/grpc/sandbox.rs | 2 +- crates/openshell-server/src/lib.rs | 4 +- crates/openshell-server/src/multiplex.rs | 9 +- crates/openshell-server/tests/common/mod.rs | 18 +- .../tests/supervisor_relay_integration.rs | 10 +- proto/openshell.proto | 30 +++ 28 files changed, 388 insertions(+), 129 deletions(-) diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index 55d7e7a29d..a072ac0d23 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -203,6 +203,15 @@ identity is absent, malformed, unknown, ambiguous, or resolves to UID/GID 0. The supervisor itself remains root so it can establish isolation before starting unprivileged children. +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. + 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..defb0019c2 100644 --- a/architecture/gateway.md +++ b/architecture/gateway.md @@ -155,24 +155,28 @@ 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, 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 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 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 7ed148304c..7d927fa209 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 5fa2c97029..f0d416bc31 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..c27ff4dbb7 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, + 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) @@ -295,7 +295,7 @@ async fn acquire_k8s_sandbox_token( .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 +308,23 @@ 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("IssueSandboxToken bootstrap exchange failed")?; - Ok(resp.into_inner().token) + .wrap_err("RegisterSupervisorPod bootstrap stream failed")?; + let mut stream = resp.into_inner(); + let activation = stream + .message() + .await + .into_diagnostic() + .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.token) } /// Build an authenticated channel for direct external use (e.g. the diff --git a/crates/openshell-core/src/sandbox_env.rs b/crates/openshell-core/src/sandbox_env.rs index 1549258fa3..dc90f8a183 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-driver-kubernetes/README.md b/crates/openshell-driver-kubernetes/README.md index 1356e2d932..a3ad2166e8 100644 --- a/crates/openshell-driver-kubernetes/README.md +++ b/crates/openshell-driver-kubernetes/README.md @@ -55,8 +55,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/config.rs b/crates/openshell-driver-kubernetes/src/config.rs index fb471180a9..9d13eeb14b 100644 --- a/crates/openshell-driver-kubernetes/src/config.rs +++ b/crates/openshell-driver-kubernetes/src/config.rs @@ -278,9 +278,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 diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index 2d947b8a28..8757cb5a8a 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -2189,7 +2189,7 @@ struct SandboxPodParams<'a> { workspace_storage_class: &'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, @@ -2419,7 +2419,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. @@ -2627,9 +2627,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, @@ -2969,7 +2970,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, diff --git a/crates/openshell-driver-kubernetes/src/main.rs b/crates/openshell-driver-kubernetes/src/main.rs index c7b0939888..5e1fc15584 100644 --- a/crates/openshell-driver-kubernetes/src/main.rs +++ b/crates/openshell-driver-kubernetes/src/main.rs @@ -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, 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/src/auth/authenticator.rs b/crates/openshell-server/src/auth/authenticator.rs index f5d5c7b2af..bc52729224 100644 --- a/crates/openshell-server/src/auth/authenticator.rs +++ b/crates/openshell-server/src/auth/authenticator.rs @@ -15,7 +15,7 @@ //! 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`) +//! tokens (path-scoped to Kubernetes 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/k8s_sa.rs b/crates/openshell-server/src/auth/k8s_sa.rs index eed0e5f083..c1b578f4b3 100644 --- a/crates/openshell-server/src/auth/k8s_sa.rs +++ b/crates/openshell-server/src/auth/k8s_sa.rs @@ -3,13 +3,13 @@ //! Kubernetes `ServiceAccount` bootstrap authenticator. //! -//! Path-scoped to `IssueSandboxToken`. Validates a projected SA token +//! Path-scoped to the Kubernetes supervisor bootstrap RPCs. 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 +//! [`SandboxIdentitySource::K8sServiceAccount`]. The bootstrap handlers then +//! mint 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 @@ -30,9 +30,13 @@ use std::sync::Arc; use tonic::Status; use tracing::{debug, info, 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 (return `Ok(None)`) so a gateway-minted JWT +/// is required there. pub const ISSUE_SANDBOX_TOKEN_PATH: &str = "/openshell.v1.OpenShell/IssueSandboxToken"; +/// Supervisor pod registration path accepted by this authenticator for the +/// phase-1 cold Kubernetes activation flow. +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 @@ -95,9 +99,9 @@ impl Authenticator for K8sServiceAccountAuthenticator { 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 { + // Scope: only Kubernetes bootstrap RPCs. Other paths fall through so + // the SandboxJwtAuthenticator (or OIDC) handles them. + if path != ISSUE_SANDBOX_TOKEN_PATH && path != REGISTER_SUPERVISOR_POD_PATH { return Ok(None); } @@ -881,7 +885,7 @@ mod tests { } #[tokio::test] - async fn authenticates_on_issue_path_only() { + async fn authenticates_on_bootstrap_paths_only() { let resolved = ResolvedK8sIdentity { sandbox_id: "sandbox-a".to_string(), pod_name: "openshell-sandbox-a".to_string(), @@ -906,6 +910,22 @@ mod tests { _ => panic!("expected sandbox principal"), } + let on_register = auth + .authenticate(&bearer_headers("sa-jwt"), REGISTER_SUPERVISOR_POD_PATH) + .await + .unwrap() + .expect("expected principal"); + match on_register { + Principal::Sandbox(p) => { + assert_eq!(p.sandbox_id, "sandbox-a"); + assert!(matches!( + p.source, + SandboxIdentitySource::K8sServiceAccount { .. } + )); + } + _ => panic!("expected sandbox principal"), + } + let off_issue = auth .authenticate( &bearer_headers("sa-jwt"), @@ -915,11 +935,11 @@ mod tests { .unwrap(); assert!( off_issue.is_none(), - "K8s SA authenticator must be scoped to IssueSandboxToken" + "K8s SA authenticator must be scoped to bootstrap paths" ); assert_eq!( fake.seen_tokens.lock().unwrap().len(), - 1, + 2, "off-path call must not consult the apiserver" ); } diff --git a/crates/openshell-server/src/auth/principal.rs b/crates/openshell-server/src/auth/principal.rs index 1d4cb7276c..e32f01c2e0 100644 --- a/crates/openshell-server/src/auth/principal.rs +++ b/crates/openshell-server/src/auth/principal.rs @@ -70,7 +70,7 @@ 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. + /// K8s `ServiceAccount` token used to bootstrap a gateway-minted JWT. + /// Populated only on Kubernetes bootstrap RPC paths. K8sServiceAccount { pod_name: String, pod_uid: 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..a0d07ce1df 100644 --- a/crates/openshell-server/src/auth/sandbox_methods.rs +++ b/crates/openshell-server/src/auth/sandbox_methods.rs @@ -26,6 +26,9 @@ mod tests { "/openshell.v1.OpenShell/ConnectSupervisor" )); assert!(is_sandbox_callable("/openshell.v1.OpenShell/RelayStream")); + assert!(is_sandbox_callable( + "/openshell.v1.OpenShell/RegisterSupervisorPod" + )); assert!(is_sandbox_callable( "/openshell.v1.OpenShell/GetSandboxConfig" )); 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..b23f8ab132 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,20 +13,55 @@ //! 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 openshell_core::proto::{ - IssueSandboxTokenRequest, IssueSandboxTokenResponse, RefreshSandboxTokenRequest, - RefreshSandboxTokenResponse, Sandbox, + IssueSandboxTokenRequest, IssueSandboxTokenResponse, PodActivationMessage, + RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, Sandbox, }; 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> { + // 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_cold_pod_activation(state, &sandbox, "IssueSandboxToken").await?; + Ok(Response::new(IssueSandboxTokenResponse { + token: activation.token, + expires_at_ms: activation.token_expires_at_ms, + })) +} + +#[allow(clippy::result_large_err, clippy::unused_async)] +pub async fn handle_register_supervisor_pod( + state: &Arc, + request: Request, +) -> Result, Status> { + let sandbox = require_k8s_bootstrap_sandbox(request.extensions(), "RegisterSupervisorPod")?; + let activation = mint_cold_pod_activation(state, &sandbox, "RegisterSupervisorPod").await?; + info!( + sandbox_id = %activation.sandbox_id, + "activated registered supervisor pod" + ); + Ok(Response::new(Box::pin(tokio_stream::once(Ok(activation))))) +} + +#[allow(clippy::result_large_err, clippy::unused_async)] +pub async fn handle_refresh_sandbox_token( + state: &Arc, + request: Request, +) -> Result, Status> { let principal = request .extensions() .get::() @@ -34,99 +70,112 @@ pub async fn handle_issue_sandbox_token( let Principal::Sandbox(sandbox) = principal else { return Err(Status::permission_denied( - "IssueSandboxToken requires a sandbox principal", + "RefreshSandboxToken requires a sandbox principal", )); }; - // 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 { .. } - ) { + // 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, - "IssueSandboxToken rejected: non-bootstrap principal source" + "RefreshSandboxToken rejected: non-gateway-JWT principal source" ); return Err(Status::permission_denied( - "this principal cannot mint a sandbox token; use RefreshSandboxToken", + "this principal cannot refresh; use RegisterSupervisorPod for bootstrap", )); - } + }; 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" + "RefreshSandboxToken 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?; + load_sandbox(state, &sandbox.sandbox_id).await?; let minted = issuer.mint(&sandbox.sandbox_id)?; info!( sandbox_id = %sandbox.sandbox_id, - "issued gateway sandbox JWT" + "renewed gateway sandbox JWT" ); - Ok(Response::new(IssueSandboxTokenResponse { + + Ok(Response::new(RefreshSandboxTokenResponse { token: minted.token, expires_at_ms: minted.expires_at_ms, })) } -#[allow(clippy::result_large_err, clippy::unused_async)] -pub async fn handle_refresh_sandbox_token( - state: &Arc, - request: Request, -) -> Result, Status> { - let principal = request - .extensions() +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( - "RefreshSandboxToken requires a sandbox principal", - )); + return Err(Status::permission_denied(format!( + "{rpc_name} requires a sandbox principal" + ))); }; - // Only callers already holding a gateway-minted JWT may refresh; the - // K8s bootstrap path must use `IssueSandboxToken`. - let SandboxIdentitySource::BootstrapJwt { .. } = &sandbox.source else { + if !matches!( + sandbox.source, + SandboxIdentitySource::K8sServiceAccount { .. } + ) { debug!( sandbox_id = %sandbox.sandbox_id, - "RefreshSandboxToken rejected: non-gateway-JWT principal source" + rpc = rpc_name, + "bootstrap RPC rejected: non-K8s ServiceAccount principal source" ); return Err(Status::permission_denied( - "this principal cannot refresh; use IssueSandboxToken for bootstrap", + "this principal cannot mint a sandbox token; use RefreshSandboxToken", )); - }; + } + Ok(sandbox) +} + +async fn mint_cold_pod_activation( + state: &Arc, + sandbox: &SandboxPrincipal, + rpc_name: &'static str, +) -> Result { let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { warn!( sandbox_id = %sandbox.sandbox_id, - "RefreshSandboxToken called but sandbox JWT issuer is not configured" + rpc = rpc_name, + "bootstrap RPC 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 record = load_sandbox(state, &sandbox.sandbox_id).await?; let minted = issuer.mint(&sandbox.sandbox_id)?; + let sandbox_name = record + .metadata + .as_ref() + .map_or_else(String::new, |m| m.name.clone()); info!( sandbox_id = %sandbox.sandbox_id, - "renewed gateway sandbox JWT" + rpc = rpc_name, + "issued gateway sandbox JWT for cold pod activation" ); - Ok(Response::new(RefreshSandboxTokenResponse { + Ok(PodActivationMessage { + sandbox_id: sandbox.sandbox_id.clone(), + sandbox_name, token: minted.token, - expires_at_ms: minted.expires_at_ms, - })) + token_expires_at_ms: minted.expires_at_ms, + startup_metadata: std::collections::HashMap::default(), + }) } -async fn ensure_sandbox_exists(state: &Arc, sandbox_id: &str) -> Result<(), Status> { +async fn load_sandbox(state: &Arc, sandbox_id: &str) -> Result { if sandbox_id.is_empty() { return Err(Status::invalid_argument("sandbox_id is required")); } @@ -136,9 +185,7 @@ async fn ensure_sandbox_exists(state: &Arc, sandbox_id: &str) -> Re .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(()) + .ok_or_else(|| Status::not_found("sandbox not found")) } #[cfg(test)] @@ -159,6 +206,7 @@ mod tests { use openshell_core::proto::{Sandbox, SandboxPhase, SandboxSpec}; 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"); @@ -273,6 +321,39 @@ mod tests { assert!(resp.expires_at_ms > 0); } + #[tokio::test] + async fn register_supervisor_pod_returns_immediate_activation_for_existing_sandbox() { + use crate::auth::principal::SandboxIdentitySource; + + let state = state_with_issuer().await; + let mut req = Request::new(RegisterSupervisorPodRequest {}); + 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(), + }, + trust_domain: Some("openshell".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 issue_rejects_missing_sandbox() { use crate::auth::principal::SandboxIdentitySource; @@ -316,9 +397,9 @@ 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. + // K8s SA-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 {}); diff --git a/crates/openshell-server/src/grpc/mod.rs b/crates/openshell-server/src/grpc/mod.rs index eed96c3598..f59e642b52 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 = "sandbox")] + 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/sandbox.rs b/crates/openshell-server/src/grpc/sandbox.rs index 0d6a4e277f..4b321e7887 100644 --- a/crates/openshell-server/src/grpc/sandbox.rs +++ b/crates/openshell-server/src/grpc/sandbox.rs @@ -220,7 +220,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/lib.rs b/crates/openshell-server/src/lib.rs index f2967a833d..d41847f389 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -136,7 +136,7 @@ pub struct ServerState { 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>, @@ -147,7 +147,7 @@ pub struct ServerState { pub sandbox_jwt_authenticator: Option>, /// Optional K8s `ServiceAccount` authenticator that backs the - /// `IssueSandboxToken` bootstrap path. Only present when the gateway + /// Kubernetes supervisor bootstrap paths. Only present when the gateway /// runs in-cluster. pub k8s_sa_authenticator: Option>, diff --git a/crates/openshell-server/src/multiplex.rs b/crates/openshell-server/src/multiplex.rs index a58f66c916..dd6f083051 100644 --- a/crates/openshell-server/src/multiplex.rs +++ b/crates/openshell-server/src/multiplex.rs @@ -694,10 +694,10 @@ 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. `K8sServiceAccountAuthenticator` (path-scoped to Kubernetes bootstrap +/// RPCs) — resolves a projected SA token to a `Principal::Sandbox` so the +/// bootstrap handler can mint a gateway JWT. No-op on every other path; +/// only present when the gateway runs in-cluster. /// 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 @@ -2131,6 +2131,7 @@ mod tests { "/openshell.v1.OpenShell/ConnectSupervisor", "/openshell.v1.OpenShell/RelayStream", "/openshell.v1.OpenShell/IssueSandboxToken", + "/openshell.v1.OpenShell/RegisterSupervisorPod", "/openshell.v1.OpenShell/RefreshSandboxToken", "/openshell.inference.v1.Inference/GetInferenceBundle", ] { 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/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). From 1b91b34a38c61837ae9def6ce4f42df5efa2e39a Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Mon, 20 Jul 2026 19:15:56 +0100 Subject: [PATCH 03/15] feat(server): add warm supervisor pod registration groundwork Signed-off-by: Gordon Sim --- crates/openshell-server-macros/src/lib.rs | 11 +- crates/openshell-server/src/auth/guard.rs | 7 +- crates/openshell-server/src/auth/k8s_sa.rs | 175 ++++++----- .../openshell-server/src/auth/method_authz.rs | 15 +- crates/openshell-server/src/auth/principal.rs | 27 ++ .../src/auth/sandbox_methods.rs | 6 +- crates/openshell-server/src/grpc/auth_rpc.rs | 103 +++++-- crates/openshell-server/src/grpc/mod.rs | 2 +- crates/openshell-server/src/grpc/policy.rs | 3 + crates/openshell-server/src/inference.rs | 5 +- crates/openshell-server/src/lib.rs | 9 + crates/openshell-server/src/multiplex.rs | 82 +++++- .../src/supervisor_pod_registration.rs | 276 ++++++++++++++++++ 13 files changed, 608 insertions(+), 113 deletions(-) create mode 100644 crates/openshell-server/src/supervisor_pod_registration.rs 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/guard.rs b/crates/openshell-server/src/auth/guard.rs index edcd6bc013..481cfed2ff 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::K8sPod`] is rejected — pod 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::K8sPod(_) => 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::K8sPod(_) => 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 c1b578f4b3..d288ebed91 100644 --- a/crates/openshell-server/src/auth/k8s_sa.rs +++ b/crates/openshell-server/src/auth/k8s_sa.rs @@ -6,17 +6,16 @@ //! Path-scoped to the Kubernetes supervisor bootstrap RPCs. 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 bootstrap handlers then -//! mint 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`]. +//! and returns either a sandbox principal for the legacy token RPC or a +//! registration-scoped pod principal for `RegisterSupervisorPod`. Warm pods may +//! register before they are bound to a sandbox, so registration identity must +//! not grant sandbox-scoped RPC access. //! //! This is the only authenticator that talks to the K8s apiserver. It is //! optional — the gateway boots without it in singleplayer deployments. use super::authenticator::Authenticator; -use super::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; +use super::principal::{Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal}; use async_trait::async_trait; use k8s_openapi::api::{ authentication::v1::{TokenReview, TokenReviewSpec, TokenReviewStatus, UserInfo}, @@ -54,24 +53,16 @@ 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>; + /// 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>; } /// Authenticator wrapper around a [`K8sIdentityResolver`]. @@ -118,18 +109,20 @@ impl Authenticator for K8sServiceAccountAuthenticator { return Ok(None); }; - if resolved.sandbox_id.is_empty() { + if path == REGISTER_SUPERVISOR_POD_PATH { + return Ok(Some(Principal::K8sPod(resolved))); + } + + let sandbox_id = resolved.sandbox_id.clone().ok_or_else(|| { warn!( pod = %resolved.pod_name, - "pod missing openshell.io/sandbox-id annotation; rejecting" + "pod missing openshell.io/sandbox-id annotation; rejecting legacy token issue" ); - return Err(Status::permission_denied( - "pod is not bound to a sandbox identity", - )); - } + Status::permission_denied("pod is not bound to a sandbox identity") + })?; Ok(Some(Principal::Sandbox(SandboxPrincipal { - sandbox_id: resolved.sandbox_id, + sandbox_id, source: SandboxIdentitySource::K8sServiceAccount { pod_name: resolved.pod_name, pod_uid: resolved.pod_uid, @@ -222,7 +215,7 @@ impl LiveK8sResolver { #[async_trait] impl K8sIdentityResolver for LiveK8sResolver { - async fn resolve(&self, token: &str) -> Result, Status> { + async fn resolve(&self, token: &str) -> Result, Status> { let review = TokenReview { metadata: ObjectMeta::default(), spec: TokenReviewSpec { @@ -295,7 +288,7 @@ impl K8sIdentityResolver for LiveK8sResolver { return Err(Status::permission_denied("SA token pod UID mismatch")); } - let sandbox_id = pod_sandbox_id(&pod)?; + 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| { @@ -317,12 +310,17 @@ impl K8sIdentityResolver for LiveK8sResolver { ); return Err(Status::permission_denied("sandbox owner not found")); }; - validate_sandbox_owner_reference(&owner, &sandbox_id, &sandbox_cr)?; + validate_sandbox_owner_reference(&owner, &sandbox_cr)?; + if let Some(ref sandbox_id) = sandbox_id { + validate_sandbox_owner_binding(&owner, sandbox_id, &sandbox_cr)?; + } - Ok(Some(ResolvedK8sIdentity { + Ok(Some(RegisteredPodIdentity { sandbox_id, pod_name: identity.pod_name, pod_uid: identity.pod_uid, + sandbox_owner_name: owner.name, + sandbox_owner_uid: owner.uid, })) } } @@ -393,20 +391,13 @@ fn user_extra_one(user: &UserInfo, key: &str) -> Result { } #[allow(clippy::result_large_err)] -fn pod_sandbox_id(pod: &Pod) -> Result { - let sandbox_id = pod - .metadata +fn pod_sandbox_id(pod: &Pod) -> Option { + 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) + .filter(|sandbox_id| !sandbox_id.is_empty()) } #[allow(clippy::result_large_err)] @@ -478,7 +469,6 @@ fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { #[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(); @@ -492,6 +482,15 @@ fn validate_sandbox_owner_reference( 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 @@ -521,12 +520,12 @@ pub mod test_support { /// 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 outcome: Result, Status>, pub seen_tokens: Mutex>, } impl FakeResolver { - pub fn returning(outcome: Result, Status>) -> Self { + pub fn returning(outcome: Result, Status>) -> Self { Self { outcome, seen_tokens: Mutex::new(Vec::new()), @@ -536,7 +535,7 @@ pub mod test_support { #[async_trait] impl K8sIdentityResolver for FakeResolver { - async fn resolve(&self, token: &str) -> Result, Status> { + 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()), @@ -655,6 +654,16 @@ mod tests { } } + fn registered_pod(sandbox_id: Option<&str>) -> RegisteredPodIdentity { + RegisteredPodIdentity { + sandbox_id: sandbox_id.map(str::to_string), + pod_name: "openshell-sandbox-a".to_string(), + pod_uid: "uid-a".to_string(), + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "cr-uid-a".to_string(), + } + } + 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); @@ -769,15 +778,12 @@ mod tests { } #[test] - fn pod_sandbox_id_requires_annotation() { + fn pod_sandbox_id_is_optional_for_warm_registration() { assert_eq!( - pod_sandbox_id(&pod_with_sandbox_id(Some("sandbox-id-a"))).unwrap(), - "sandbox-id-a" + pod_sandbox_id(&pod_with_sandbox_id(Some("sandbox-id-a"))).as_deref(), + Some("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); + assert!(pod_sandbox_id(&pod_with_sandbox_id(None)).is_none()); } #[test] @@ -863,34 +869,40 @@ mod tests { } #[test] - fn validate_sandbox_owner_reference_requires_matching_cr_uid_and_label() { + 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, "sandbox-id-a", &cr) - .expect("matching CR should be accepted"); + 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, "sandbox-id-a", &wrong_uid) + 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_reference(&owner, "sandbox-id-a", &wrong_label) + 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 authenticates_on_bootstrap_paths_only() { - let resolved = ResolvedK8sIdentity { - sandbox_id: "sandbox-a".to_string(), - pod_name: "openshell-sandbox-a".to_string(), - pod_uid: "uid-a".to_string(), - }; + let resolved = registered_pod(Some("sandbox-a")); let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); let auth = K8sServiceAccountAuthenticator::new(fake.clone()); @@ -916,14 +928,12 @@ mod tests { .unwrap() .expect("expected principal"); match on_register { - Principal::Sandbox(p) => { - assert_eq!(p.sandbox_id, "sandbox-a"); - assert!(matches!( - p.source, - SandboxIdentitySource::K8sServiceAccount { .. } - )); + Principal::K8sPod(p) => { + assert_eq!(p.sandbox_id.as_deref(), Some("sandbox-a")); + assert_eq!(p.pod_name, "openshell-sandbox-a"); + assert_eq!(p.pod_uid, "uid-a"); } - _ => panic!("expected sandbox principal"), + _ => panic!("expected pod registration principal"), } let off_issue = auth @@ -970,12 +980,8 @@ mod tests { } #[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(), - }; + async fn issue_token_rejects_unbound_registered_pod() { + let resolved = registered_pod(None); let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); let auth = K8sServiceAccountAuthenticator::new(fake); let err = auth @@ -985,6 +991,27 @@ mod tests { assert_eq!(err.code(), tonic::Code::PermissionDenied); } + #[tokio::test] + async fn register_supervisor_pod_accepts_unbound_registered_pod() { + let resolved = registered_pod(None); + let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); + let auth = K8sServiceAccountAuthenticator::new(fake); + let principal = auth + .authenticate(&bearer_headers("sa-jwt"), REGISTER_SUPERVISOR_POD_PATH) + .await + .expect("unbound warm pod registration should authenticate") + .expect("expected principal"); + + match principal { + Principal::K8sPod(p) => { + assert!(p.sandbox_id.is_none()); + assert_eq!(p.pod_name, "openshell-sandbox-a"); + assert_eq!(p.sandbox_owner_uid, "cr-uid-a"); + } + _ => panic!("expected pod registration principal"), + } + } + #[tokio::test] async fn resolver_error_propagates() { let fake = Arc::new(FakeResolver::returning(Err(Status::unavailable( 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 e32f01c2e0..d395805250 100644 --- a/crates/openshell-server/src/auth/principal.rs +++ b/crates/openshell-server/src/auth/principal.rs @@ -28,6 +28,12 @@ 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), + /// Kubernetes pod authenticated for supervisor pod registration only. + /// + /// This is intentionally not a sandbox principal: warm pods can be valid + /// Kubernetes pods before they are claimed by a sandbox. The router only + /// allows this principal to call `RegisterSupervisorPod`. + K8sPod(RegisteredPodIdentity), /// Truly unauthenticated caller (health probes, reflection). Sandbox-class /// and user-class methods reject this variant. #[allow(dead_code)] @@ -41,6 +47,27 @@ pub struct UserPrincipal { pub identity: Identity, } +/// Kubernetes pod identity validated by `TokenReview` and live pod lookup. +/// +/// `sandbox_id` is present only when the pod is already bound to a concrete +/// `OpenShell` sandbox. Unbound warm pods carry the owning Sandbox CR metadata +/// but must not receive a gateway sandbox JWT until a later claim binds the +/// exact pod UID to a sandbox. +#[derive(Debug, Clone)] +#[allow(dead_code)] +pub struct RegisteredPodIdentity { + /// Live pod name from the `TokenReview` binding. + pub pod_name: String, + /// Live pod UID from the `TokenReview` binding and pod lookup. + pub pod_uid: String, + /// `OpenShell` sandbox UUID from the pod annotation, when already bound. + pub sandbox_id: Option, + /// Owning Sandbox CR name from the pod ownerReference. + pub sandbox_owner_name: String, + /// Owning Sandbox CR UID from the pod ownerReference and live CR lookup. + pub sandbox_owner_uid: String, +} + /// Sandbox caller — bound to one specific sandbox UUID. /// /// `sandbox_id` and `source` are consumed by the router and handler guards. diff --git a/crates/openshell-server/src/auth/sandbox_methods.rs b/crates/openshell-server/src/auth/sandbox_methods.rs index a0d07ce1df..2566c29ab5 100644 --- a/crates/openshell-server/src/auth/sandbox_methods.rs +++ b/crates/openshell-server/src/auth/sandbox_methods.rs @@ -26,9 +26,6 @@ mod tests { "/openshell.v1.OpenShell/ConnectSupervisor" )); assert!(is_sandbox_callable("/openshell.v1.OpenShell/RelayStream")); - assert!(is_sandbox_callable( - "/openshell.v1.OpenShell/RegisterSupervisorPod" - )); assert!(is_sandbox_callable( "/openshell.v1.OpenShell/GetSandboxConfig" )); @@ -54,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/grpc/auth_rpc.rs b/crates/openshell-server/src/grpc/auth_rpc.rs index b23f8ab132..52d26a0eba 100644 --- a/crates/openshell-server/src/grpc/auth_rpc.rs +++ b/crates/openshell-server/src/grpc/auth_rpc.rs @@ -13,7 +13,9 @@ //! until their own `exp` and are bounded by the configured short TTL. use crate::ServerState; -use crate::auth::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; +use crate::auth::principal::{ + Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, +}; use openshell_core::proto::{ IssueSandboxTokenRequest, IssueSandboxTokenResponse, PodActivationMessage, RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, Sandbox, @@ -36,7 +38,8 @@ pub async fn handle_issue_sandbox_token( // 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_cold_pod_activation(state, &sandbox, "IssueSandboxToken").await?; + let activation = + mint_cold_pod_activation(state, &sandbox.sandbox_id, "IssueSandboxToken").await?; Ok(Response::new(IssueSandboxTokenResponse { token: activation.token, expires_at_ms: activation.token_expires_at_ms, @@ -48,13 +51,21 @@ pub async fn handle_register_supervisor_pod( state: &Arc, request: Request, ) -> Result, Status> { - let sandbox = require_k8s_bootstrap_sandbox(request.extensions(), "RegisterSupervisorPod")?; - let activation = mint_cold_pod_activation(state, &sandbox, "RegisterSupervisorPod").await?; - info!( - sandbox_id = %activation.sandbox_id, - "activated registered supervisor pod" - ); - Ok(Response::new(Box::pin(tokio_stream::once(Ok(activation))))) + let pod = require_registered_pod(request.extensions())?; + if let Some(sandbox_id) = pod.sandbox_id.as_deref() { + let activation = + mint_cold_pod_activation(state, sandbox_id, "RegisterSupervisorPod").await?; + info!( + sandbox_id = %activation.sandbox_id, + pod = %pod.pod_name, + pod_uid = %pod.pod_uid, + "activated bound supervisor pod" + ); + return Ok(Response::new(Box::pin(tokio_stream::once(Ok(activation))))); + } + + let stream = state.supervisor_pod_registrations.register_pending(pod)?; + Ok(Response::new(Box::pin(stream))) } #[allow(clippy::result_large_err, clippy::unused_async)] @@ -140,34 +151,53 @@ fn require_k8s_bootstrap_sandbox( Ok(sandbox) } +fn require_registered_pod(extensions: &tonic::Extensions) -> Result { + if let Some(pod) = extensions.get::().cloned() { + return Ok(pod); + } + + let principal = extensions + .get::() + .cloned() + .ok_or_else(|| Status::unauthenticated("missing principal"))?; + + let Principal::K8sPod(pod) = principal else { + return Err(Status::permission_denied( + "RegisterSupervisorPod requires a Kubernetes pod registration principal", + )); + }; + + Ok(pod) +} + async fn mint_cold_pod_activation( state: &Arc, - sandbox: &SandboxPrincipal, + sandbox_id: &str, rpc_name: &'static str, ) -> Result { let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { warn!( - sandbox_id = %sandbox.sandbox_id, + sandbox_id = %sandbox_id, rpc = rpc_name, "bootstrap RPC called but sandbox JWT issuer is not configured" ); Status::unavailable("sandbox JWT minting is not configured on this gateway") })?; - let record = load_sandbox(state, &sandbox.sandbox_id).await?; - let minted = issuer.mint(&sandbox.sandbox_id)?; + 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.sandbox_id, + sandbox_id = %sandbox_id, rpc = rpc_name, "issued gateway sandbox JWT for cold pod activation" ); Ok(PodActivationMessage { - sandbox_id: sandbox.sandbox_id.clone(), + sandbox_id: sandbox_id.to_string(), sandbox_name, token: minted.token, token_expires_at_ms: minted.expires_at_ms, @@ -273,6 +303,16 @@ mod tests { }) } + fn registered_pod(sandbox_id: Option<&str>) -> RegisteredPodIdentity { + RegisteredPodIdentity { + pod_name: "pod-a".to_string(), + pod_uid: "uid-a".to_string(), + sandbox_id: sandbox_id.map(str::to_string), + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "owner-uid-a".to_string(), + } + } + #[tokio::test] async fn refresh_returns_new_token() { let state = state_with_issuer().await; @@ -323,19 +363,10 @@ mod tests { #[tokio::test] async fn register_supervisor_pod_returns_immediate_activation_for_existing_sandbox() { - use crate::auth::principal::SandboxIdentitySource; - let state = state_with_issuer().await; let mut req = Request::new(RegisterSupervisorPodRequest {}); 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(), - }, - trust_domain: Some("openshell".to_string()), - })); + .insert(Principal::K8sPod(registered_pod(Some("sandbox-a")))); let mut stream = handle_register_supervisor_pod(&state, req) .await @@ -354,6 +385,28 @@ mod tests { 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::K8sPod(registered_pod(None))); + + 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; diff --git a/crates/openshell-server/src/grpc/mod.rs b/crates/openshell-server/src/grpc/mod.rs index f59e642b52..b1278adf94 100644 --- a/crates/openshell-server/src/grpc/mod.rs +++ b/crates/openshell-server/src/grpc/mod.rs @@ -692,7 +692,7 @@ impl OpenShell for OpenShellService { Box> + Send + 'static>, >; - #[rpc_auth(auth = "sandbox")] + #[rpc_auth(auth = "pod_registration")] async fn register_supervisor_pod( &self, request: Request, diff --git a/crates/openshell-server/src/grpc/policy.rs b/crates/openshell-server/src/grpc/policy.rs index e99ea1a30d..ffb766d1e7 100644 --- a/crates/openshell-server/src/grpc/policy.rs +++ b/crates/openshell-server/src/grpc/policy.rs @@ -1253,6 +1253,9 @@ async fn resolve_sandbox_by_name_for_principal( Ok(sandbox) } Principal::User(_) => sandbox.ok_or_else(|| Status::not_found("sandbox not found")), + Principal::K8sPod(_) => 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/inference.rs b/crates/openshell-server/src/inference.rs index 39d6afe2fd..d2cda430b6 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::K8sPod(_), + ) => 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 d41847f389..539d4ee7fa 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"))] @@ -132,6 +133,11 @@ 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>, @@ -202,6 +208,9 @@ 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, diff --git a/crates/openshell-server/src/multiplex.rs b/crates/openshell-server/src/multiplex.rs index dd6f083051..231fb4999b 100644 --- a/crates/openshell-server/src/multiplex.rs +++ b/crates/openshell-server/src/multiplex.rs @@ -471,6 +471,22 @@ fn gateway_principal_fields(principal: &Principal) -> BTreeMap { fields.insert("trust_domain".to_string(), trust_domain.clone()); } } + Principal::K8sPod(pod) => { + fields.insert("kind".to_string(), "k8s_pod".to_string()); + fields.insert("pod_name".to_string(), pod.pod_name.clone()); + fields.insert("pod_uid".to_string(), pod.pod_uid.clone()); + fields.insert( + "sandbox_owner_name".to_string(), + pod.sandbox_owner_name.clone(), + ); + fields.insert( + "sandbox_owner_uid".to_string(), + pod.sandbox_owner_uid.clone(), + ); + if let Some(sandbox_id) = &pod.sandbox_id { + fields.insert("sandbox_id".to_string(), sandbox_id.clone()); + } + } Principal::Anonymous => { fields.insert("kind".to_string(), "anonymous".to_string()); } @@ -695,8 +711,9 @@ where /// /// Chain order (first-match-wins): /// 1. `K8sServiceAccountAuthenticator` (path-scoped to Kubernetes bootstrap -/// RPCs) — resolves a projected SA token to a `Principal::Sandbox` so the -/// bootstrap handler can mint a gateway JWT. No-op on every other path; +/// RPCs) — resolves a projected SA token to either a legacy +/// `Principal::Sandbox` for `IssueSandboxToken` or a registration-scoped +/// pod principal for `RegisterSupervisorPod`. No-op on every other path; /// only present when the gateway runs in-cluster. /// 2. `SandboxJwtAuthenticator` — validates gateway-minted JWTs. Recognized /// via a distinctive `kid` so non-matching Bearer tokens fall through. @@ -743,6 +760,7 @@ 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::K8sPod` is gated by the pod-registration method only. #[derive(Clone)] pub struct AuthGrpcRouter { inner: S, @@ -881,6 +899,14 @@ where ))); } } + Principal::K8sPod(ref pod) => { + if !crate::auth::method_authz::is_pod_registration_callable(&path) { + return Ok(status_response(tonic::Status::permission_denied( + "Kubernetes pod principals may only register supervisor pods", + ))); + } + req.extensions_mut().insert(pod.clone()); + } Principal::Anonymous => { return Ok(status_response(tonic::Status::unauthenticated( "anonymous callers may not call authenticated methods", @@ -1821,7 +1847,8 @@ mod tests { use crate::auth::authenticator::test_support::MockAuthenticator; use crate::auth::identity::{Identity, IdentityProvider}; use crate::auth::principal::{ - Principal, SandboxIdentitySource, SandboxPrincipal, UserPrincipal, + Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, + UserPrincipal, }; use http_body_util::Full; use std::sync::Arc; @@ -1911,6 +1938,16 @@ mod tests { }) } + fn pod_registration_principal() -> Principal { + Principal::K8sPod(RegisteredPodIdentity { + pod_name: "pod-a".to_string(), + pod_uid: "pod-uid-a".to_string(), + sandbox_id: None, + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "owner-uid-a".to_string(), + }) + } + #[tokio::test] async fn mtls_peer_identity_fills_missing_principal_when_enabled() { let mock = Arc::new(MockAuthenticator::returning(Ok(None))); @@ -2099,6 +2136,43 @@ mod tests { )); } + #[tokio::test] + async fn pod_registration_principal_can_only_register_supervisor_pod() { + let mock = Arc::new(MockAuthenticator::returning(Ok(Some( + pod_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::K8sPod(_)) + )); + + let mock = Arc::new(MockAuthenticator::returning(Ok(Some( + pod_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 @@ -2131,7 +2205,6 @@ mod tests { "/openshell.v1.OpenShell/ConnectSupervisor", "/openshell.v1.OpenShell/RelayStream", "/openshell.v1.OpenShell/IssueSandboxToken", - "/openshell.v1.OpenShell/RegisterSupervisorPod", "/openshell.v1.OpenShell/RefreshSandboxToken", "/openshell.inference.v1.Inference/GetInferenceBundle", ] { @@ -2156,6 +2229,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/supervisor_pod_registration.rs b/crates/openshell-server/src/supervisor_pod_registration.rs new file mode 100644 index 0000000000..1c847c43d4 --- /dev/null +++ b/crates/openshell-server/src/supervisor_pod_registration.rs @@ -0,0 +1,276 @@ +// 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 crate::auth::principal::RegisteredPodIdentity; +use openshell_core::proto::PodActivationMessage; +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_pod_uid: HashMap, +} + +#[derive(Debug)] +struct PendingRegistration { + identity: RegisteredPodIdentity, + 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: RegisteredPodIdentity, + ) -> Result { + if identity.pod_uid.is_empty() { + return Err(Status::permission_denied("registered pod UID is required")); + } + + let (sender, receiver) = mpsc::channel(1); + let session_id = self.next_session_id.fetch_add(1, Ordering::Relaxed); + let pod_uid = identity.pod_uid.clone(); + let pod_name = identity.pod_name.clone(); + let sandbox_owner = identity.sandbox_owner_name.clone(); + + let replaced = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + inner + .pending_by_pod_uid + .insert( + pod_uid.clone(), + PendingRegistration { + identity, + sender, + session_id, + registered_at: Instant::now(), + }, + ) + .is_some() + }; + + if replaced { + info!( + pod = %pod_name, + pod_uid = %pod_uid, + sandbox_owner = %sandbox_owner, + "replaced duplicate pending supervisor pod registration" + ); + } else { + info!( + pod = %pod_name, + pod_uid = %pod_uid, + sandbox_owner = %sandbox_owner, + "registered warm supervisor pod pending activation" + ); + } + + Ok(PendingRegistrationStream { + registry: self.clone(), + pod_uid, + session_id, + inner: ReceiverStream::new(receiver), + }) + } + + #[allow(clippy::result_large_err)] + #[allow(dead_code)] + pub fn activate( + &self, + pod_uid: &str, + activation: PodActivationMessage, + ) -> Result { + let pending = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + inner.pending_by_pod_uid.remove(pod_uid) + }; + + let Some(pending) = pending else { + debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to activate"); + return Ok(false); + }; + + let pod_name = pending.identity.pod_name.clone(); + pending.sender.try_send(Ok(activation)).map_err(|_| { + warn!( + pod = %pod_name, + pod_uid = %pod_uid, + "pending supervisor pod registration stream closed before activation" + ); + Status::unavailable("registered pod stream closed before activation") + })?; + info!( + pod = %pod_name, + pod_uid = %pod_uid, + pending_ms = pending.registered_at.elapsed().as_millis(), + "activated pending supervisor pod registration" + ); + Ok(true) + } + + #[must_use] + #[allow(dead_code)] + pub fn pending_count(&self) -> usize { + self.inner + .lock() + .expect("pending pod registry poisoned") + .pending_by_pod_uid + .len() + } + + fn remove_if_session(&self, pod_uid: &str, session_id: u64) { + let removed = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner + .pending_by_pod_uid + .get(pod_uid) + .is_some_and(|pending| pending.session_id == session_id) + { + inner.pending_by_pod_uid.remove(pod_uid) + } else { + None + } + }; + + if let Some(pending) = removed { + debug!( + pod = %pending.identity.pod_name, + pod_uid = %pod_uid, + pending_ms = pending.registered_at.elapsed().as_millis(), + "removed pending supervisor pod registration" + ); + } + } +} + +pub struct PendingRegistrationStream { + registry: Arc, + pod_uid: 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.pod_uid, self.session_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + use tokio_stream::StreamExt; + + fn pod_identity(pod_uid: &str) -> RegisteredPodIdentity { + RegisteredPodIdentity { + pod_name: "warm-pod-a".to_string(), + pod_uid: pod_uid.to_string(), + sandbox_id: None, + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "owner-uid-a".to_string(), + } + } + + #[test] + fn dropping_stream_removes_pending_registration() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let stream = registry + .register_pending(pod_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(pod_identity("pod-uid-a")) + .expect("register old"); + let new = registry + .register_pending(pod_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(pod_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(), + }; + assert!( + registry + .activate("pod-uid-a", activation) + .expect("activate") + ); + assert_eq!(registry.pending_count(), 0); + + let received = stream + .next() + .await + .expect("activation message") + .expect("activation OK"); + assert_eq!(received.sandbox_id, "sandbox-a"); + assert!(stream.next().await.is_none()); + } +} From 514169f2d2ef9f0bb8c1f9715316919cb6bc92ee Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Mon, 20 Jul 2026 19:58:56 +0100 Subject: [PATCH 04/15] feat(server): scaffold claim-driven warm pod activation Signed-off-by: Gordon Sim --- crates/openshell-server/src/grpc/auth_rpc.rs | 57 +-- crates/openshell-server/src/lib.rs | 1 + .../src/supervisor_pod_registration.rs | 180 ++++++++- .../src/warm_pod_activation.rs | 381 ++++++++++++++++++ 4 files changed, 551 insertions(+), 68 deletions(-) create mode 100644 crates/openshell-server/src/warm_pod_activation.rs diff --git a/crates/openshell-server/src/grpc/auth_rpc.rs b/crates/openshell-server/src/grpc/auth_rpc.rs index 52d26a0eba..7bfdcc820d 100644 --- a/crates/openshell-server/src/grpc/auth_rpc.rs +++ b/crates/openshell-server/src/grpc/auth_rpc.rs @@ -16,9 +16,10 @@ use crate::ServerState; use crate::auth::principal::{ Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, }; +use crate::warm_pod_activation::{load_sandbox, mint_pod_activation}; use openshell_core::proto::{ IssueSandboxTokenRequest, IssueSandboxTokenResponse, PodActivationMessage, - RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, Sandbox, + RefreshSandboxTokenRequest, RefreshSandboxTokenResponse, RegisterSupervisorPodRequest, }; use std::sync::Arc; use std::{pin::Pin, result::Result as StdResult}; @@ -38,8 +39,7 @@ pub async fn handle_issue_sandbox_token( // 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_cold_pod_activation(state, &sandbox.sandbox_id, "IssueSandboxToken").await?; + 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, @@ -53,8 +53,7 @@ pub async fn handle_register_supervisor_pod( ) -> Result, Status> { let pod = require_registered_pod(request.extensions())?; if let Some(sandbox_id) = pod.sandbox_id.as_deref() { - let activation = - mint_cold_pod_activation(state, sandbox_id, "RegisterSupervisorPod").await?; + let activation = mint_pod_activation(state, sandbox_id, "RegisterSupervisorPod").await?; info!( sandbox_id = %activation.sandbox_id, pod = %pod.pod_name, @@ -170,54 +169,6 @@ fn require_registered_pod(extensions: &tonic::Extensions) -> Result, - sandbox_id: &str, - rpc_name: &'static str, -) -> Result { - let issuer = state.sandbox_jwt_issuer.as_ref().ok_or_else(|| { - warn!( - sandbox_id = %sandbox_id, - rpc = rpc_name, - "bootstrap RPC called 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, - rpc = rpc_name, - "issued gateway sandbox JWT for cold 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(), - }) -} - -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::*; diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 539d4ee7fa..88ebcd6d89 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -52,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; diff --git a/crates/openshell-server/src/supervisor_pod_registration.rs b/crates/openshell-server/src/supervisor_pod_registration.rs index 1c847c43d4..a8eb9eafa0 100644 --- a/crates/openshell-server/src/supervisor_pod_registration.rs +++ b/crates/openshell-server/src/supervisor_pod_registration.rs @@ -32,6 +32,7 @@ pub struct SupervisorPodRegistrationRegistry { #[derive(Debug, Default)] struct Inner { pending_by_pod_uid: HashMap, + activated_pod_uid: HashMap, } #[derive(Debug)] @@ -65,6 +66,11 @@ impl SupervisorPodRegistrationRegistry { let replaced = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_pod_uid.contains_key(&pod_uid) { + return Err(Status::already_exists( + "supervisor pod has already been activated", + )); + } inner .pending_by_pod_uid .insert( @@ -104,38 +110,103 @@ impl SupervisorPodRegistrationRegistry { } #[allow(clippy::result_large_err)] - #[allow(dead_code)] - pub fn activate( + pub(crate) fn pending_identity(&self, pod_uid: &str) -> Result { + let inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_pod_uid.contains_key(pod_uid) { + return Err(Status::already_exists( + "supervisor pod has already been activated", + )); + } + inner + .pending_by_pod_uid + .get(pod_uid) + .map(|pending| pending.identity.clone()) + .ok_or_else(|| Status::not_found("pending supervisor pod registration not found")) + } + + #[allow(clippy::result_large_err)] + pub(crate) fn activate( &self, pod_uid: &str, activation: PodActivationMessage, - ) -> Result { + ) -> Result<(), Status> { + let pending = { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_pod_uid.contains_key(pod_uid) { + return Err(Status::already_exists( + "supervisor pod has already been activated", + )); + } + let Some(pending) = inner.pending_by_pod_uid.remove(pod_uid) else { + debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to activate"); + return Err(Status::not_found( + "pending supervisor pod registration not found", + )); + }; + inner + .activated_pod_uid + .insert(pod_uid.to_string(), Instant::now()); + pending + }; + + let pod_name = pending.identity.pod_name.clone(); + if pending.sender.try_send(Ok(activation)).is_err() { + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + inner.activated_pod_uid.remove(pod_uid); + warn!( + pod = %pod_name, + pod_uid = %pod_uid, + "pending supervisor pod registration stream closed before activation" + ); + return Err(Status::unavailable( + "registered pod stream closed before activation", + )); + } + + info!( + pod = %pod_name, + pod_uid = %pod_uid, + pending_ms = pending.registered_at.elapsed().as_millis(), + "activated pending supervisor pod registration" + ); + Ok(()) + } + + #[allow(clippy::result_large_err)] + pub(crate) fn fail_pending(&self, pod_uid: &str, status: Status) -> Result<(), Status> { let pending = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + if inner.activated_pod_uid.contains_key(pod_uid) { + return Err(Status::already_exists( + "supervisor pod has already been activated", + )); + } inner.pending_by_pod_uid.remove(pod_uid) }; let Some(pending) = pending else { - debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to activate"); - return Ok(false); + debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to fail"); + return Err(Status::not_found( + "pending supervisor pod registration not found", + )); }; let pod_name = pending.identity.pod_name.clone(); - pending.sender.try_send(Ok(activation)).map_err(|_| { + pending.sender.try_send(Err(status)).map_err(|_| { warn!( pod = %pod_name, pod_uid = %pod_uid, - "pending supervisor pod registration stream closed before activation" + "pending supervisor pod registration stream closed before failure notification" ); - Status::unavailable("registered pod stream closed before activation") + Status::unavailable("registered pod stream closed before failure notification") })?; info!( pod = %pod_name, pod_uid = %pod_uid, pending_ms = pending.registered_at.elapsed().as_millis(), - "activated pending supervisor pod registration" + "failed pending supervisor pod registration" ); - Ok(true) + Ok(()) } #[must_use] @@ -148,6 +219,16 @@ impl SupervisorPodRegistrationRegistry { .len() } + #[must_use] + #[allow(dead_code)] + pub fn activated_count(&self) -> usize { + self.inner + .lock() + .expect("pending pod registry poisoned") + .activated_pod_uid + .len() + } + fn remove_if_session(&self, pod_uid: &str, session_id: u64) { let removed = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); @@ -258,12 +339,11 @@ mod tests { token_expires_at_ms: 123, startup_metadata: HashMap::default(), }; - assert!( - registry - .activate("pod-uid-a", activation) - .expect("activate") - ); + registry + .activate("pod-uid-a", activation) + .expect("activate"); assert_eq!(registry.pending_count(), 0); + assert_eq!(registry.activated_count(), 1); let received = stream .next() @@ -273,4 +353,74 @@ mod tests { 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(pod_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(pod_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(pod_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/warm_pod_activation.rs b/crates/openshell-server/src/warm_pod_activation.rs new file mode 100644 index 0000000000..881ff34651 --- /dev/null +++ b/crates/openshell-server/src/warm_pod_activation.rs @@ -0,0 +1,381 @@ +// 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 pod UID to an `OpenShell` sandbox. + +use crate::ServerState; +use crate::auth::principal::RegisteredPodIdentity; +use async_trait::async_trait; +use openshell_core::proto::{PodActivationMessage, Sandbox}; +use std::sync::Arc; +use tonic::Status; +use tracing::{info, warn}; + +#[derive(Debug, Clone, PartialEq, Eq)] +#[allow(dead_code)] +pub struct WarmPodActivationTarget { + pub pod_uid: String, + pub sandbox_id: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +#[allow(dead_code)] +pub struct ValidatedWarmPodActivation { + pub pod_uid: String, + pub pod_name: String, + pub sandbox_owner_uid: String, + pub sandbox_id: String, +} + +#[async_trait] +#[allow(dead_code)] +pub trait WarmPodActivationValidator: Send + Sync { + async fn validate( + &self, + target: &WarmPodActivationTarget, + pending: &RegisteredPodIdentity, + ) -> Result; +} + +#[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.pod_uid)?; + let activation_result = async { + let validated = validator.validate(&target, &pending).await?; + ensure_validated_activation_matches_pending(&target, &pending, &validated)?; + mint_pod_activation(state, &validated.sandbox_id, "WarmPodActivation").await + } + .await; + let activation = match activation_result { + Ok(activation) => activation, + Err(status) => { + let stream_status = clone_status(&status); + state + .supervisor_pod_registrations + .fail_pending(&target.pod_uid, stream_status)?; + return Err(status); + } + }; + + state + .supervisor_pod_registrations + .activate(&target.pod_uid, 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: &RegisteredPodIdentity, + validated: &ValidatedWarmPodActivation, +) -> Result<(), Status> { + if validated.pod_uid != target.pod_uid || validated.pod_uid != pending.pod_uid { + return Err(Status::permission_denied( + "validated pod UID does not match pending registration", + )); + } + if validated.pod_name != pending.pod_name { + return Err(Status::permission_denied( + "validated pod name does not match pending registration", + )); + } + if validated.sandbox_owner_uid != pending.sandbox_owner_uid { + return Err(Status::permission_denied( + "validated Sandbox 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: &RegisteredPodIdentity, + ) -> 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(), + resource_version: 0, + }), + 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_pod() -> RegisteredPodIdentity { + RegisteredPodIdentity { + pod_name: "warm-pod-a".to_string(), + pod_uid: "pod-uid-a".to_string(), + sandbox_id: None, + sandbox_owner_name: "sandbox-owner-a".to_string(), + sandbox_owner_uid: "owner-uid-a".to_string(), + } + } + + fn activation_target(sandbox_id: &str) -> WarmPodActivationTarget { + WarmPodActivationTarget { + pod_uid: "pod-uid-a".to_string(), + sandbox_id: sandbox_id.to_string(), + } + } + + fn validated_activation(sandbox_id: &str) -> ValidatedWarmPodActivation { + ValidatedWarmPodActivation { + pod_uid: "pod-uid-a".to_string(), + pod_name: "warm-pod-a".to_string(), + sandbox_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_pod()) + .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_pod_uid_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 pod UID 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_pod()) + .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(), 0); + let stream_err = stream + .next() + .await + .expect("failure message") + .expect_err("missing sandbox should fail stream"); + assert_eq!(stream_err.code(), tonic::Code::NotFound); + } + + #[tokio::test] + async fn duplicate_activation_for_same_pod_uid_is_rejected() { + let state = state_with_issuer().await; + let mut stream = state + .supervisor_pod_registrations + .register_pending(pending_pod()) + .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_pod()) + .expect("register pending"); + let mut validated = validated_activation("sandbox-a"); + validated.sandbox_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); + } +} From 16b906f6745c4bc7b59808fbd3e361bb3c359355 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Tue, 21 Jul 2026 20:31:30 +0100 Subject: [PATCH 05/15] refactor(kubernetes): move supervisor bootstrap identity into driver Signed-off-by: Gordon Sim --- crates/openshell-core/src/lib.rs | 1 + .../src/supervisor_bootstrap.rs | 94 ++ .../src/bootstrap.rs | 907 ++++++++++++++++ crates/openshell-driver-kubernetes/src/lib.rs | 4 + .../src/auth/authenticator.rs | 4 +- crates/openshell-server/src/auth/guard.rs | 8 +- crates/openshell-server/src/auth/k8s_sa.rs | 995 ++---------------- crates/openshell-server/src/auth/principal.rs | 39 +- .../src/compute/driver_config.rs | 45 - crates/openshell-server/src/compute/mod.rs | 47 + crates/openshell-server/src/grpc/auth_rpc.rs | 91 +- crates/openshell-server/src/grpc/policy.rs | 2 +- crates/openshell-server/src/inference.rs | 2 +- crates/openshell-server/src/lib.rs | 42 - crates/openshell-server/src/multiplex.rs | 79 +- .../src/supervisor_pod_registration.rs | 182 ++-- .../src/warm_pod_activation.rs | 141 ++- 17 files changed, 1474 insertions(+), 1209 deletions(-) create mode 100644 crates/openshell-core/src/supervisor_bootstrap.rs create mode 100644 crates/openshell-driver-kubernetes/src/bootstrap.rs diff --git a/crates/openshell-core/src/lib.rs b/crates/openshell-core/src/lib.rs index 80bfbb046d..ee857ea8ba 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; pub mod transport_errors; diff --git a/crates/openshell-core/src/supervisor_bootstrap.rs b/crates/openshell-core/src/supervisor_bootstrap.rs new file mode 100644 index 0000000000..3ffabb737c --- /dev/null +++ b/crates/openshell-core/src/supervisor_bootstrap.rs @@ -0,0 +1,94 @@ +// 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, + /// 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-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/lib.rs b/crates/openshell-driver-kubernetes/src/lib.rs index 7c56c8de5b..f967170e43 100644 --- a/crates/openshell-driver-kubernetes/src/lib.rs +++ b/crates/openshell-driver-kubernetes/src/lib.rs @@ -1,10 +1,14 @@ // 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 use bootstrap::{ + K8sIdentityResolver, KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, +}; pub use config::{ AppArmorProfile, DEFAULT_PROXY_UID, DEFAULT_SANDBOX_SERVICE_ACCOUNT_NAME, DEFAULT_WORKSPACE_STORAGE_SIZE, KubernetesComputeConfig, KubernetesSidecarConfig, diff --git a/crates/openshell-server/src/auth/authenticator.rs b/crates/openshell-server/src/auth/authenticator.rs index bc52729224..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 Kubernetes bootstrap RPCs) +//! - [`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 481cfed2ff..c1e87b6e48 100644 --- a/crates/openshell-server/src/auth/guard.rs +++ b/crates/openshell-server/src/auth/guard.rs @@ -20,8 +20,8 @@ use tracing::info; /// scope at the router level). /// - [`Principal::Sandbox`] must reference the same canonical UUID it /// was authenticated with. -/// - [`Principal::K8sPod`] is rejected — pod registration identity is not a -/// sandbox-scoped credential. +/// - [`Principal::SupervisorBootstrap`] is rejected — bootstrap registration +/// identity is not a sandbox-scoped credential. /// - [`Principal::Anonymous`] is rejected — sandbox-class methods are /// never anonymously callable. /// @@ -46,7 +46,7 @@ pub fn ensure_sandbox_scope(principal: &Principal, claimed_sandbox_id: &str) -> )) } } - Principal::K8sPod(_) => Err(Status::permission_denied( + Principal::SupervisorBootstrap(_) => Err(Status::permission_denied( "sandbox-scoped methods require a sandbox principal", )), Principal::Anonymous => Err(Status::unauthenticated( @@ -89,7 +89,7 @@ pub fn ensure_sandbox_principal_scope( ensure_sandbox_scope(principal, claimed_sandbox_id)?; Ok(p.clone()) } - Principal::User(_) | Principal::K8sPod(_) => 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 d288ebed91..fd547530ca 100644 --- a/crates/openshell-server/src/auth/k8s_sa.rs +++ b/crates/openshell-server/src/auth/k8s_sa.rs @@ -1,97 +1,58 @@ // 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 the Kubernetes supervisor bootstrap RPCs. 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 either a sandbox principal for the legacy token RPC or a -//! registration-scoped pod principal for `RegisterSupervisorPod`. Warm pods may -//! register before they are bound to a sandbox, so registration identity must -//! not grant sandbox-scoped RPC access. -//! -//! 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, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal}; -use async_trait::async_trait; -use k8s_openapi::api::{ - authentication::v1::{TokenReview, TokenReviewSpec, TokenReviewStatus, UserInfo}, - core::v1::Pod, +use super::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; +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}; /// Legacy gRPC method path accepted by this authenticator. All other -/// non-bootstrap paths fall through (return `Ok(None)`) so a gateway-minted JWT +/// 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 pod registration path accepted by this authenticator for the -/// phase-1 cold Kubernetes activation flow. +/// 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"; - -/// 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>; +/// Path-scoped authenticator backed by the active compute driver's bootstrap +/// identity provider. +pub struct SupervisorBootstrapAuthenticator { + provider: Arc, } -/// Authenticator wrapper around a [`K8sIdentityResolver`]. -pub struct K8sServiceAccountAuthenticator { - resolver: 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 Kubernetes bootstrap RPCs. Other paths fall through so - // the SandboxJwtAuthenticator (or OIDC) handles them. if path != ISSUE_SANDBOX_TOKEN_PATH && path != REGISTER_SUPERVISOR_POD_PATH { return Ok(None); } @@ -104,428 +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 path == REGISTER_SUPERVISOR_POD_PATH { - return Ok(Some(Principal::K8sPod(resolved))); + return Ok(Some(Principal::SupervisorBootstrap(identity))); } - let sandbox_id = resolved.sandbox_id.clone().ok_or_else(|| { + let SupervisorBootstrapBinding::BoundSandbox { sandbox_id } = identity.binding else { warn!( - pod = %resolved.pod_name, - "pod missing openshell.io/sandbox-id annotation; rejecting legacy token issue" + 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" ); - Status::permission_denied("pod is not bound to a sandbox identity") - })?; + return Err(Status::permission_denied( + "supervisor instance is not bound to a sandbox identity", + )); + }; Ok(Some(Principal::Sandbox(SandboxPrincipal { sandbox_id, - source: SandboxIdentitySource::K8sServiceAccount { - pod_name: resolved.pod_name, - pod_uid: resolved.pod_uid, + 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_cr)?; - if let Some(ref sandbox_id) = sandbox_id { - validate_sandbox_owner_binding(&owner, sandbox_id, &sandbox_cr)?; - } - - Ok(Some(RegisteredPodIdentity { - sandbox_id, - pod_name: identity.pod_name, - pod_uid: identity.pod_uid, - sandbox_owner_name: owner.name, - sandbox_owner_uid: owner.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) -> 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 { +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()), @@ -534,430 +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, - }) - } - - #[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::>(), - ), - }), - } + headers } - 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 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 registered_pod(sandbox_id: Option<&str>) -> RegisteredPodIdentity { - RegisteredPodIdentity { - sandbox_id: sandbox_id.map(str::to_string), - pod_name: "openshell-sandbox-a".to_string(), - pod_uid: "uid-a".to_string(), - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "cr-uid-a".to_string(), - } - } - - 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 authenticates_on_bootstrap_paths_only() { - let resolved = registered_pod(Some("sandbox-a")); - let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); - let auth = K8sServiceAccountAuthenticator::new(fake.clone()); + 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 on_register = auth - .authenticate(&bearer_headers("sa-jwt"), REGISTER_SUPERVISOR_POD_PATH) + let register = auth + .authenticate( + &bearer_headers("bootstrap-token"), + REGISTER_SUPERVISOR_POD_PATH, + ) .await .unwrap() .expect("expected principal"); - match on_register { - Principal::K8sPod(p) => { - assert_eq!(p.sandbox_id.as_deref(), Some("sandbox-a")); - assert_eq!(p.pod_name, "openshell-sandbox-a"); - assert_eq!(p.pod_uid, "uid-a"); - } - _ => panic!("expected pod registration principal"), - } + assert!(matches!(register, Principal::SupervisorBootstrap(_))); - let off_issue = auth + 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 bootstrap paths" - ); - assert_eq!( - fake.seen_tokens.lock().unwrap().len(), - 2, - "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 @@ -966,62 +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 issue_token_rejects_unbound_registered_pod() { - let resolved = registered_pod(None); - 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 register_supervisor_pod_accepts_unbound_registered_pod() { - let resolved = registered_pod(None); - let fake = Arc::new(FakeResolver::returning(Ok(Some(resolved)))); - let auth = K8sServiceAccountAuthenticator::new(fake); + 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("sa-jwt"), REGISTER_SUPERVISOR_POD_PATH) + .authenticate( + &bearer_headers("bootstrap-token"), + REGISTER_SUPERVISOR_POD_PATH, + ) .await - .expect("unbound warm pod registration should authenticate") + .unwrap() .expect("expected principal"); match principal { - Principal::K8sPod(p) => { - assert!(p.sandbox_id.is_none()); - assert_eq!(p.pod_name, "openshell-sandbox-a"); - assert_eq!(p.sandbox_owner_uid, "cr-uid-a"); + Principal::SupervisorBootstrap(identity) => { + assert_eq!(identity.binding, SupervisorBootstrapBinding::WarmPending); + assert_eq!(identity.owner_uid, "cr-uid-a"); } - _ => panic!("expected pod registration principal"), + _ => panic!("expected bootstrap principal"), } } - - #[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) - .await - .expect_err("resolver error must propagate"); - assert_eq!(err.code(), tonic::Code::Unavailable); - } } diff --git a/crates/openshell-server/src/auth/principal.rs b/crates/openshell-server/src/auth/principal.rs index d395805250..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,12 +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), - /// Kubernetes pod authenticated for supervisor pod registration only. + /// Driver runtime instance authenticated for supervisor bootstrap + /// registration only. /// /// This is intentionally not a sandbox principal: warm pods can be valid - /// Kubernetes pods before they are claimed by a sandbox. The router only + /// driver instances before they are claimed by a sandbox. The router only /// allows this principal to call `RegisterSupervisorPod`. - K8sPod(RegisteredPodIdentity), + SupervisorBootstrap(SupervisorBootstrapIdentity), /// Truly unauthenticated caller (health probes, reflection). Sandbox-class /// and user-class methods reject this variant. #[allow(dead_code)] @@ -47,27 +49,6 @@ pub struct UserPrincipal { pub identity: Identity, } -/// Kubernetes pod identity validated by `TokenReview` and live pod lookup. -/// -/// `sandbox_id` is present only when the pod is already bound to a concrete -/// `OpenShell` sandbox. Unbound warm pods carry the owning Sandbox CR metadata -/// but must not receive a gateway sandbox JWT until a later claim binds the -/// exact pod UID to a sandbox. -#[derive(Debug, Clone)] -#[allow(dead_code)] -pub struct RegisteredPodIdentity { - /// Live pod name from the `TokenReview` binding. - pub pod_name: String, - /// Live pod UID from the `TokenReview` binding and pod lookup. - pub pod_uid: String, - /// `OpenShell` sandbox UUID from the pod annotation, when already bound. - pub sandbox_id: Option, - /// Owning Sandbox CR name from the pod ownerReference. - pub sandbox_owner_name: String, - /// Owning Sandbox CR UID from the pod ownerReference and live CR lookup. - pub sandbox_owner_uid: String, -} - /// Sandbox caller — bound to one specific sandbox UUID. /// /// `sandbox_id` and `source` are consumed by the router and handler guards. @@ -97,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. - /// Populated only on Kubernetes bootstrap RPC paths. - 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/compute/driver_config.rs b/crates/openshell-server/src/compute/driver_config.rs index f56d233f2f..2a3e0f5795 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<'_>, @@ -265,35 +249,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 b87f1b5fa6..34d92d84ad 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -37,10 +37,12 @@ use openshell_core::proto::{ PlatformEvent, Sandbox, SandboxCondition, SandboxPhase, SandboxSpec, SandboxStatus, SandboxTemplate, ServiceEndpoint, SshSession, }; +use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentityProvider; use openshell_core::{ObjectLabels, ObjectWorkspace}; use openshell_driver_docker::DockerComputeDriver; use openshell_driver_kubernetes::{ ComputeDriverService as KubernetesDriverService, KubernetesComputeDriver, + KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, }; use openshell_driver_podman::{ComputeDriverService as PodmanDriverService, PodmanComputeDriver}; use prost::Message; @@ -231,6 +233,33 @@ 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 + } + } +} + #[derive(Debug)] pub struct AcquiredRemoteDriverEndpoint { pub(crate) name: String, @@ -368,6 +397,7 @@ pub struct ComputeRuntime { sandbox_watch_bus: SandboxWatchBus, tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, + supervisor_bootstrap_identity: Option>, sync_lock: Arc>, delete_gates: Arc, gateway_bind_addresses: Vec, @@ -393,6 +423,7 @@ impl ComputeRuntime { sandbox_watch_bus: SandboxWatchBus, tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, + supervisor_bootstrap_identity: Option>, gateway_bind_addresses: Vec, ) -> Result { let capabilities = driver @@ -425,6 +456,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + supervisor_bootstrap_identity, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses, @@ -486,6 +518,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + None, gateway_bind_addresses, ) .await @@ -499,6 +532,8 @@ impl ComputeRuntime { tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, ) -> Result { + let supervisor_bootstrap_identity = + kubernetes_supervisor_bootstrap_identity_provider(&config).await; let driver = KubernetesComputeDriver::new(config) .await .map_err(|err| ComputeError::Message(err.to_string()))?; @@ -514,6 +549,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + supervisor_bootstrap_identity, Vec::new(), ) .await @@ -539,6 +575,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + None, Vec::new(), ) .await @@ -567,6 +604,7 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, + None, Vec::new(), ) .await @@ -582,6 +620,13 @@ impl ComputeRuntime { std::slice::from_ref(&self.driver_info) } + #[must_use] + pub fn supervisor_bootstrap_identity_provider( + &self, + ) -> Option> { + self.supervisor_bootstrap_identity.clone() + } + #[must_use] pub fn driver_kind(&self) -> Option { self.driver_info.name.parse().ok() @@ -2845,6 +2890,7 @@ pub async fn new_test_runtime_for_driver(store: Arc, driver_name: &str) - sandbox_watch_bus: SandboxWatchBus::new(), tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), + supervisor_bootstrap_identity: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses: Vec::new(), @@ -3314,6 +3360,7 @@ mod tests { sandbox_watch_bus: SandboxWatchBus::new(), tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), + supervisor_bootstrap_identity: 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/grpc/auth_rpc.rs b/crates/openshell-server/src/grpc/auth_rpc.rs index 7bfdcc820d..84c0a47a66 100644 --- a/crates/openshell-server/src/grpc/auth_rpc.rs +++ b/crates/openshell-server/src/grpc/auth_rpc.rs @@ -13,14 +13,13 @@ //! until their own `exp` and are bounded by the configured short TTL. use crate::ServerState; -use crate::auth::principal::{ - Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, -}; +use crate::auth::principal::{Principal, SandboxIdentitySource, SandboxPrincipal}; use crate::warm_pod_activation::{load_sandbox, mint_pod_activation}; use openshell_core::proto::{ 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; @@ -51,19 +50,22 @@ pub async fn handle_register_supervisor_pod( state: &Arc, request: Request, ) -> Result, Status> { - let pod = require_registered_pod(request.extensions())?; - if let Some(sandbox_id) = pod.sandbox_id.as_deref() { + 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, - pod = %pod.pod_name, - pod_uid = %pod.pod_uid, - "activated bound supervisor pod" + driver = %identity.driver, + instance_name = %identity.instance_name, + instance_id = %identity.instance_id, + "activated bound supervisor instance" ); return Ok(Response::new(Box::pin(tokio_stream::once(Ok(activation))))); } - let stream = state.supervisor_pod_registrations.register_pending(pod)?; + let stream = state + .supervisor_pod_registrations + .register_pending(identity)?; Ok(Response::new(Box::pin(stream))) } @@ -135,12 +137,12 @@ fn require_k8s_bootstrap_sandbox( if !matches!( sandbox.source, - SandboxIdentitySource::K8sServiceAccount { .. } + SandboxIdentitySource::SupervisorBootstrap { .. } ) { debug!( sandbox_id = %sandbox.sandbox_id, rpc = rpc_name, - "bootstrap RPC rejected: non-K8s ServiceAccount principal source" + "bootstrap RPC rejected: non-bootstrap principal source" ); return Err(Status::permission_denied( "this principal cannot mint a sandbox token; use RefreshSandboxToken", @@ -150,9 +152,11 @@ fn require_k8s_bootstrap_sandbox( Ok(sandbox) } -fn require_registered_pod(extensions: &tonic::Extensions) -> Result { - if let Some(pod) = extensions.get::().cloned() { - return Ok(pod); +fn require_bootstrap_identity( + extensions: &tonic::Extensions, +) -> Result { + if let Some(identity) = extensions.get::().cloned() { + return Ok(identity); } let principal = extensions @@ -160,13 +164,13 @@ fn require_registered_pod(extensions: &tonic::Extensions) -> Result) -> RegisteredPodIdentity { - RegisteredPodIdentity { - pod_name: "pod-a".to_string(), - pod_uid: "uid-a".to_string(), - sandbox_id: sandbox_id.map(str::to_string), - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + 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, } } @@ -298,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()), })); @@ -317,7 +324,11 @@ mod tests { let state = state_with_issuer().await; let mut req = Request::new(RegisterSupervisorPodRequest {}); req.extensions_mut() - .insert(Principal::K8sPod(registered_pod(Some("sandbox-a")))); + .insert(Principal::SupervisorBootstrap(bootstrap_identity( + SupervisorBootstrapBinding::BoundSandbox { + sandbox_id: "sandbox-a".to_string(), + }, + ))); let mut stream = handle_register_supervisor_pod(&state, req) .await @@ -341,7 +352,9 @@ mod tests { let state = state_with_issuer().await; let mut req = Request::new(RegisterSupervisorPodRequest {}); req.extensions_mut() - .insert(Principal::K8sPod(registered_pod(None))); + .insert(Principal::SupervisorBootstrap(bootstrap_identity( + SupervisorBootstrapBinding::WarmPending, + ))); let mut stream = handle_register_supervisor_pod(&state, req) .await @@ -367,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()), })); @@ -400,8 +414,8 @@ mod tests { } #[tokio::test] - async fn refresh_rejects_k8s_sa_principal() { - // K8s SA-bootstrap principals must use RegisterSupervisorPod, not + 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; @@ -410,15 +424,16 @@ 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()), })); 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/policy.rs b/crates/openshell-server/src/grpc/policy.rs index ffb766d1e7..10646b885c 100644 --- a/crates/openshell-server/src/grpc/policy.rs +++ b/crates/openshell-server/src/grpc/policy.rs @@ -1253,7 +1253,7 @@ async fn resolve_sandbox_by_name_for_principal( Ok(sandbox) } Principal::User(_) => sandbox.ok_or_else(|| Status::not_found("sandbox not found")), - Principal::K8sPod(_) => Err(Status::permission_denied( + Principal::SupervisorBootstrap(_) => Err(Status::permission_denied( "sandbox-scoped methods require a sandbox principal", )), Principal::Anonymous => Err(Status::unauthenticated( diff --git a/crates/openshell-server/src/inference.rs b/crates/openshell-server/src/inference.rs index d2cda430b6..e68753135e 100644 --- a/crates/openshell-server/src/inference.rs +++ b/crates/openshell-server/src/inference.rs @@ -932,7 +932,7 @@ fn authorize_inference_bundle( Some(crate::auth::principal::Principal::Sandbox(s)) => Ok(s.sandbox_id.clone()), Some( crate::auth::principal::Principal::User(_) - | crate::auth::principal::Principal::K8sPod(_), + | crate::auth::principal::Principal::SupervisorBootstrap(_), ) => Err(Status::permission_denied( "GetInferenceBundle requires a sandbox principal", )), diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 88ebcd6d89..a8c9f19d95 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -153,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 - /// Kubernetes supervisor bootstrap paths. 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, @@ -215,7 +210,6 @@ impl ServerState { oidc_cache, sandbox_jwt_issuer: None, sandbox_jwt_authenticator: None, - k8s_sa_authenticator: None, grpc_rate_limiter, gateway_interceptors: None, provider_profile_sources: @@ -390,42 +384,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); diff --git a/crates/openshell-server/src/multiplex.rs b/crates/openshell-server/src/multiplex.rs index 231fb4999b..0fa11daaff 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,21 +471,13 @@ fn gateway_principal_fields(principal: &Principal) -> BTreeMap { fields.insert("trust_domain".to_string(), trust_domain.clone()); } } - Principal::K8sPod(pod) => { - fields.insert("kind".to_string(), "k8s_pod".to_string()); - fields.insert("pod_name".to_string(), pod.pod_name.clone()); - fields.insert("pod_uid".to_string(), pod.pod_uid.clone()); - fields.insert( - "sandbox_owner_name".to_string(), - pod.sandbox_owner_name.clone(), - ); - fields.insert( - "sandbox_owner_uid".to_string(), - pod.sandbox_owner_uid.clone(), - ); - if let Some(sandbox_id) = &pod.sandbox_id { - fields.insert("sandbox_id".to_string(), sandbox_id.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()); @@ -710,11 +702,12 @@ where /// Assemble the authenticator chain for the gateway. /// /// Chain order (first-match-wins): -/// 1. `K8sServiceAccountAuthenticator` (path-scoped to Kubernetes bootstrap -/// RPCs) — resolves a projected SA token to either a legacy -/// `Principal::Sandbox` for `IssueSandboxToken` or a registration-scoped -/// pod principal for `RegisterSupervisorPod`. 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 @@ -732,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); @@ -760,7 +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::K8sPod` is gated by the pod-registration method only. +/// `Principal::SupervisorBootstrap` is gated by the pod-registration method +/// only. #[derive(Clone)] pub struct AuthGrpcRouter { inner: S, @@ -899,13 +895,13 @@ where ))); } } - Principal::K8sPod(ref pod) => { + Principal::SupervisorBootstrap(ref identity) => { if !crate::auth::method_authz::is_pod_registration_callable(&path) { return Ok(status_response(tonic::Status::permission_denied( - "Kubernetes pod principals may only register supervisor pods", + "supervisor bootstrap principals may only register supervisors", ))); } - req.extensions_mut().insert(pod.clone()); + req.extensions_mut().insert(identity.clone()); } Principal::Anonymous => { return Ok(status_response(tonic::Status::unauthenticated( @@ -1847,10 +1843,12 @@ mod tests { use crate::auth::authenticator::test_support::MockAuthenticator; use crate::auth::identity::{Identity, IdentityProvider}; use crate::auth::principal::{ - Principal, RegisteredPodIdentity, SandboxIdentitySource, SandboxPrincipal, - UserPrincipal, + 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; @@ -1938,13 +1936,14 @@ mod tests { }) } - fn pod_registration_principal() -> Principal { - Principal::K8sPod(RegisteredPodIdentity { - pod_name: "pod-a".to_string(), - pod_uid: "pod-uid-a".to_string(), - sandbox_id: None, - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + 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, }) } @@ -2137,9 +2136,9 @@ mod tests { } #[tokio::test] - async fn pod_registration_principal_can_only_register_supervisor_pod() { + async fn bootstrap_registration_principal_can_only_register_supervisor_pod() { let mock = Arc::new(MockAuthenticator::returning(Ok(Some( - pod_registration_principal(), + bootstrap_registration_principal(), )))); let chain = AuthenticatorChain::new(vec![mock]); let (recorder, seen) = PrincipalRecorder::new(); @@ -2155,11 +2154,11 @@ mod tests { assert_eq!(res.status(), 200); assert!(matches!( seen.lock().unwrap().as_ref(), - Some(Principal::K8sPod(_)) + Some(Principal::SupervisorBootstrap(_)) )); let mock = Arc::new(MockAuthenticator::returning(Ok(Some( - pod_registration_principal(), + bootstrap_registration_principal(), )))); let chain = AuthenticatorChain::new(vec![mock]); let (recorder, seen) = PrincipalRecorder::new(); diff --git a/crates/openshell-server/src/supervisor_pod_registration.rs b/crates/openshell-server/src/supervisor_pod_registration.rs index a8eb9eafa0..805cfc27a2 100644 --- a/crates/openshell-server/src/supervisor_pod_registration.rs +++ b/crates/openshell-server/src/supervisor_pod_registration.rs @@ -8,8 +8,8 @@ //! assignment; the future claim controller will activate the stored stream once //! it binds the exact pod UID to a sandbox. -use crate::auth::principal::RegisteredPodIdentity; 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; @@ -31,13 +31,13 @@ pub struct SupervisorPodRegistrationRegistry { #[derive(Debug, Default)] struct Inner { - pending_by_pod_uid: HashMap, - activated_pod_uid: HashMap, + pending_by_instance_id: HashMap, + activated_instance_id: HashMap, } #[derive(Debug)] struct PendingRegistration { - identity: RegisteredPodIdentity, + identity: SupervisorBootstrapIdentity, sender: mpsc::Sender>, session_id: u64, registered_at: Instant, @@ -52,29 +52,32 @@ impl SupervisorPodRegistrationRegistry { #[allow(clippy::result_large_err)] pub fn register_pending( self: &Arc, - identity: RegisteredPodIdentity, + identity: SupervisorBootstrapIdentity, ) -> Result { - if identity.pod_uid.is_empty() { - return Err(Status::permission_denied("registered pod UID is required")); + 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 pod_uid = identity.pod_uid.clone(); - let pod_name = identity.pod_name.clone(); - let sandbox_owner = identity.sandbox_owner_name.clone(); + 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_pod_uid.contains_key(&pod_uid) { + if inner.activated_instance_id.contains_key(&instance_id) { return Err(Status::already_exists( - "supervisor pod has already been activated", + "supervisor instance has already been activated", )); } inner - .pending_by_pod_uid + .pending_by_instance_id .insert( - pod_uid.clone(), + instance_id.clone(), PendingRegistration { identity, sender, @@ -87,124 +90,132 @@ impl SupervisorPodRegistrationRegistry { if replaced { info!( - pod = %pod_name, - pod_uid = %pod_uid, - sandbox_owner = %sandbox_owner, - "replaced duplicate pending supervisor pod registration" + driver = %driver, + instance = %instance_name, + instance_id = %instance_id, + owner = %owner_name, + "replaced duplicate pending supervisor registration" ); } else { info!( - pod = %pod_name, - pod_uid = %pod_uid, - sandbox_owner = %sandbox_owner, - "registered warm supervisor pod pending activation" + driver = %driver, + instance = %instance_name, + instance_id = %instance_id, + owner = %owner_name, + "registered warm supervisor instance pending activation" ); } Ok(PendingRegistrationStream { registry: self.clone(), - pod_uid, + instance_id, session_id, inner: ReceiverStream::new(receiver), }) } #[allow(clippy::result_large_err)] - pub(crate) fn pending_identity(&self, pod_uid: &str) -> Result { + pub(crate) fn pending_identity( + &self, + instance_id: &str, + ) -> Result { let inner = self.inner.lock().expect("pending pod registry poisoned"); - if inner.activated_pod_uid.contains_key(pod_uid) { + if inner.activated_instance_id.contains_key(instance_id) { return Err(Status::already_exists( - "supervisor pod has already been activated", + "supervisor instance has already been activated", )); } inner - .pending_by_pod_uid - .get(pod_uid) + .pending_by_instance_id + .get(instance_id) .map(|pending| pending.identity.clone()) - .ok_or_else(|| Status::not_found("pending supervisor pod registration not found")) + .ok_or_else(|| Status::not_found("pending supervisor registration not found")) } #[allow(clippy::result_large_err)] pub(crate) fn activate( &self, - pod_uid: &str, + instance_id: &str, activation: PodActivationMessage, ) -> Result<(), Status> { let pending = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); - if inner.activated_pod_uid.contains_key(pod_uid) { + if inner.activated_instance_id.contains_key(instance_id) { return Err(Status::already_exists( - "supervisor pod has already been activated", + "supervisor instance has already been activated", )); } - let Some(pending) = inner.pending_by_pod_uid.remove(pod_uid) else { - debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to activate"); + 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 pod registration not found", + "pending supervisor registration not found", )); }; inner - .activated_pod_uid - .insert(pod_uid.to_string(), Instant::now()); + .activated_instance_id + .insert(instance_id.to_string(), Instant::now()); pending }; - let pod_name = pending.identity.pod_name.clone(); + 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_pod_uid.remove(pod_uid); + inner.activated_instance_id.remove(instance_id); warn!( - pod = %pod_name, - pod_uid = %pod_uid, - "pending supervisor pod registration stream closed before activation" + instance = %instance_name, + instance_id = %instance_id, + "pending supervisor registration stream closed before activation" ); return Err(Status::unavailable( - "registered pod stream closed before activation", + "registered supervisor stream closed before activation", )); } info!( - pod = %pod_name, - pod_uid = %pod_uid, + instance = %instance_name, + instance_id = %instance_id, pending_ms = pending.registered_at.elapsed().as_millis(), - "activated pending supervisor pod registration" + "activated pending supervisor registration" ); Ok(()) } #[allow(clippy::result_large_err)] - pub(crate) fn fail_pending(&self, pod_uid: &str, status: Status) -> Result<(), Status> { + 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_pod_uid.contains_key(pod_uid) { + if inner.activated_instance_id.contains_key(instance_id) { return Err(Status::already_exists( - "supervisor pod has already been activated", + "supervisor instance has already been activated", )); } - inner.pending_by_pod_uid.remove(pod_uid) + inner.pending_by_instance_id.remove(instance_id) }; let Some(pending) = pending else { - debug!(pod_uid = %pod_uid, "no pending supervisor pod registration to fail"); + debug!(instance_id = %instance_id, "no pending supervisor registration to fail"); return Err(Status::not_found( - "pending supervisor pod registration not found", + "pending supervisor registration not found", )); }; - let pod_name = pending.identity.pod_name.clone(); + let instance_name = pending.identity.instance_name.clone(); pending.sender.try_send(Err(status)).map_err(|_| { warn!( - pod = %pod_name, - pod_uid = %pod_uid, - "pending supervisor pod registration stream closed before failure notification" + instance = %instance_name, + instance_id = %instance_id, + "pending supervisor registration stream closed before failure notification" ); - Status::unavailable("registered pod stream closed before failure notification") + Status::unavailable("registered supervisor stream closed before failure notification") })?; info!( - pod = %pod_name, - pod_uid = %pod_uid, + instance = %instance_name, + instance_id = %instance_id, pending_ms = pending.registered_at.elapsed().as_millis(), - "failed pending supervisor pod registration" + "failed pending supervisor registration" ); Ok(()) } @@ -215,7 +226,7 @@ impl SupervisorPodRegistrationRegistry { self.inner .lock() .expect("pending pod registry poisoned") - .pending_by_pod_uid + .pending_by_instance_id .len() } @@ -225,19 +236,19 @@ impl SupervisorPodRegistrationRegistry { self.inner .lock() .expect("pending pod registry poisoned") - .activated_pod_uid + .activated_instance_id .len() } - fn remove_if_session(&self, pod_uid: &str, session_id: u64) { + 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_pod_uid - .get(pod_uid) + .pending_by_instance_id + .get(instance_id) .is_some_and(|pending| pending.session_id == session_id) { - inner.pending_by_pod_uid.remove(pod_uid) + inner.pending_by_instance_id.remove(instance_id) } else { None } @@ -245,10 +256,10 @@ impl SupervisorPodRegistrationRegistry { if let Some(pending) = removed { debug!( - pod = %pending.identity.pod_name, - pod_uid = %pod_uid, + instance = %pending.identity.instance_name, + instance_id = %instance_id, pending_ms = pending.registered_at.elapsed().as_millis(), - "removed pending supervisor pod registration" + "removed pending supervisor registration" ); } } @@ -256,7 +267,7 @@ impl SupervisorPodRegistrationRegistry { pub struct PendingRegistrationStream { registry: Arc, - pod_uid: String, + instance_id: String, session_id: u64, inner: ReceiverStream>, } @@ -273,7 +284,7 @@ impl Stream for PendingRegistrationStream { impl Drop for PendingRegistrationStream { fn drop(&mut self) { self.registry - .remove_if_session(&self.pod_uid, self.session_id); + .remove_if_session(&self.instance_id, self.session_id); } } @@ -283,13 +294,14 @@ mod tests { use std::collections::HashMap; use tokio_stream::StreamExt; - fn pod_identity(pod_uid: &str) -> RegisteredPodIdentity { - RegisteredPodIdentity { - pod_name: "warm-pod-a".to_string(), - pod_uid: pod_uid.to_string(), - sandbox_id: None, - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + 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, } } @@ -297,7 +309,7 @@ mod tests { fn dropping_stream_removes_pending_registration() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let stream = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register pending"); assert_eq!(registry.pending_count(), 1); drop(stream); @@ -308,10 +320,10 @@ mod tests { fn duplicate_registration_replaces_prior_stream() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let old = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register old"); let new = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register new"); assert_eq!(registry.pending_count(), 1); @@ -329,7 +341,7 @@ mod tests { async fn activation_sends_message_and_removes_pending_registration() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let mut stream = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register pending"); let activation = PodActivationMessage { @@ -376,7 +388,7 @@ mod tests { async fn activation_tombstone_rejects_duplicate_activation_and_registration() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let mut stream = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register pending"); let activation = PodActivationMessage { @@ -396,7 +408,7 @@ mod tests { .expect_err("duplicate activation must fail"); assert_eq!(err.code(), tonic::Code::AlreadyExists); - let Err(err) = registry.register_pending(pod_identity("pod-uid-a")) else { + 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); @@ -406,7 +418,7 @@ mod tests { fn closed_pending_stream_cannot_be_activated() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); let stream = registry - .register_pending(pod_identity("pod-uid-a")) + .register_pending(bootstrap_identity("pod-uid-a")) .expect("register pending"); drop(stream); diff --git a/crates/openshell-server/src/warm_pod_activation.rs b/crates/openshell-server/src/warm_pod_activation.rs index 881ff34651..b93c9307ae 100644 --- a/crates/openshell-server/src/warm_pod_activation.rs +++ b/crates/openshell-server/src/warm_pod_activation.rs @@ -4,12 +4,15 @@ //! 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 pod UID to an `OpenShell` sandbox. +//! revalidates a claim that binds a warm supervisor instance to an `OpenShell` +//! sandbox. use crate::ServerState; -use crate::auth::principal::RegisteredPodIdentity; 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}; @@ -17,16 +20,19 @@ use tracing::{info, warn}; #[derive(Debug, Clone, PartialEq, Eq)] #[allow(dead_code)] pub struct WarmPodActivationTarget { - pub pod_uid: String, + 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 pod_uid: String, - pub pod_name: String, - pub sandbox_owner_uid: String, + pub driver: String, + pub instance_id: String, + pub instance_name: String, + pub owner_uid: String, pub sandbox_id: String, } @@ -36,10 +42,59 @@ pub trait WarmPodActivationValidator: Send + Sync { async fn validate( &self, target: &WarmPodActivationTarget, - pending: &RegisteredPodIdentity, + 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, + }; + activate_warm_pod(&self.state, &TrustDriverActivationRequest, target).await + } +} + +struct TrustDriverActivationRequest; + +#[async_trait] +impl WarmPodActivationValidator for TrustDriverActivationRequest { + async fn validate( + &self, + target: &WarmPodActivationTarget, + pending: &SupervisorBootstrapIdentity, + ) -> Result { + 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( @@ -52,7 +107,7 @@ where { let pending = state .supervisor_pod_registrations - .pending_identity(&target.pod_uid)?; + .pending_identity(&target.instance_id)?; let activation_result = async { let validated = validator.validate(&target, &pending).await?; ensure_validated_activation_matches_pending(&target, &pending, &validated)?; @@ -65,14 +120,14 @@ where let stream_status = clone_status(&status); state .supervisor_pod_registrations - .fail_pending(&target.pod_uid, stream_status)?; + .fail_pending(&target.instance_id, stream_status)?; return Err(status); } }; state .supervisor_pod_registrations - .activate(&target.pod_uid, activation)?; + .activate(&target.instance_id, activation)?; Ok(()) } @@ -83,22 +138,27 @@ fn clone_status(status: &Status) -> Status { #[allow(clippy::result_large_err)] fn ensure_validated_activation_matches_pending( target: &WarmPodActivationTarget, - pending: &RegisteredPodIdentity, + pending: &SupervisorBootstrapIdentity, validated: &ValidatedWarmPodActivation, ) -> Result<(), Status> { - if validated.pod_uid != target.pod_uid || validated.pod_uid != pending.pod_uid { + 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 pod UID does not match pending registration", + "validated instance ID does not match pending registration", )); } - if validated.pod_name != pending.pod_name { + if validated.instance_name != pending.instance_name { return Err(Status::permission_denied( - "validated pod name does not match pending registration", + "validated instance name does not match pending registration", )); } - if validated.sandbox_owner_uid != pending.sandbox_owner_uid { + if validated.owner_uid != target.owner_uid || validated.owner_uid != pending.owner_uid { return Err(Status::permission_denied( - "validated Sandbox owner UID does not match pending registration", + "validated owner UID does not match pending registration", )); } if validated.sandbox_id != target.sandbox_id { @@ -185,7 +245,7 @@ mod tests { async fn validate( &self, _target: &WarmPodActivationTarget, - _pending: &RegisteredPodIdentity, + _pending: &SupervisorBootstrapIdentity, ) -> Result { Ok(self.validated.clone()) } @@ -229,7 +289,10 @@ mod tests { 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, @@ -241,28 +304,32 @@ mod tests { state.store.put_message(&sandbox).await.unwrap(); } - fn pending_pod() -> RegisteredPodIdentity { - RegisteredPodIdentity { - pod_name: "warm-pod-a".to_string(), - pod_uid: "pod-uid-a".to_string(), - sandbox_id: None, - sandbox_owner_name: "sandbox-owner-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + 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 { - pod_uid: "pod-uid-a".to_string(), + 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 { - pod_uid: "pod-uid-a".to_string(), - pod_name: "warm-pod-a".to_string(), - sandbox_owner_uid: "owner-uid-a".to_string(), + 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(), } } @@ -272,7 +339,7 @@ mod tests { let state = state_with_issuer().await; let mut stream = state .supervisor_pod_registrations - .register_pending(pending_pod()) + .register_pending(pending_identity()) .expect("register pending"); let validator = StaticValidator { validated: validated_activation("sandbox-a"), @@ -295,7 +362,7 @@ mod tests { } #[tokio::test] - async fn activation_for_unknown_pod_uid_fails_before_validation() { + async fn activation_for_unknown_instance_id_fails_before_validation() { let state = state_with_issuer().await; let validator = StaticValidator { validated: validated_activation("sandbox-a"), @@ -303,7 +370,7 @@ mod tests { let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) .await - .expect_err("unknown pod UID must fail"); + .expect_err("unknown instance ID must fail"); assert_eq!(err.code(), tonic::Code::NotFound); } @@ -313,7 +380,7 @@ mod tests { let state = state_with_issuer().await; let mut stream = state .supervisor_pod_registrations - .register_pending(pending_pod()) + .register_pending(pending_identity()) .expect("register pending"); let validator = StaticValidator { validated: validated_activation("sandbox-deleted"), @@ -334,11 +401,11 @@ mod tests { } #[tokio::test] - async fn duplicate_activation_for_same_pod_uid_is_rejected() { + 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_pod()) + .register_pending(pending_identity()) .expect("register pending"); let validator = StaticValidator { validated: validated_activation("sandbox-a"), @@ -360,10 +427,10 @@ mod tests { let state = state_with_issuer().await; let mut stream = state .supervisor_pod_registrations - .register_pending(pending_pod()) + .register_pending(pending_identity()) .expect("register pending"); let mut validated = validated_activation("sandbox-a"); - validated.sandbox_owner_uid = "other-owner".to_string(); + validated.owner_uid = "other-owner".to_string(); let validator = StaticValidator { validated }; let err = activate_warm_pod(&state, &validator, activation_target("sandbox-a")) From 2f30b6646331309d5333a59216442c4ab8ac52de Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Wed, 22 Jul 2026 09:31:12 +0100 Subject: [PATCH 06/15] feat(kubernetes): activate warm pods from SandboxClaims Signed-off-by: Gordon Sim --- crates/openshell-driver-kubernetes/src/lib.rs | 2 + .../src/sandboxclaim.rs | 747 ++++++++++++++++++ crates/openshell-server/src/compute/mod.rs | 47 +- crates/openshell-server/src/lib.rs | 4 + 4 files changed, 799 insertions(+), 1 deletion(-) create mode 100644 crates/openshell-driver-kubernetes/src/sandboxclaim.rs diff --git a/crates/openshell-driver-kubernetes/src/lib.rs b/crates/openshell-driver-kubernetes/src/lib.rs index f967170e43..8c077fc7e3 100644 --- a/crates/openshell-driver-kubernetes/src/lib.rs +++ b/crates/openshell-driver-kubernetes/src/lib.rs @@ -5,6 +5,7 @@ pub mod bootstrap; pub mod config; pub mod driver; pub mod grpc; +pub mod sandboxclaim; pub use bootstrap::{ K8sIdentityResolver, KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, @@ -16,3 +17,4 @@ pub use config::{ }; pub use driver::{KubernetesComputeDriver, KubernetesDriverError}; pub use grpc::ComputeDriverService; +pub use sandboxclaim::SandboxClaimActivationController; diff --git a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs new file mode 100644 index 0000000000..833dea9bab --- /dev/null +++ b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs @@ -0,0 +1,747 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Kubernetes `SandboxClaim` activation support. + +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_SANDBOX_ID; +use openshell_core::supervisor_bootstrap::{ + SupervisorBootstrapActivationRequest, SupervisorBootstrapActivator, +}; +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 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_CLAIM_GROUP: &str = "extensions.agents.x-k8s.io"; +const SANDBOX_CLAIM_KIND: &str = "SandboxClaim"; +const SANDBOX_CLAIM_VERSION_V1BETA1: &str = "v1beta1"; +const SANDBOX_CLAIM_VERSION_V1ALPHA1: &str = "v1alpha1"; +const SANDBOX_CLAIM_VERSIONS: &[&str] = &[ + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_VERSION_V1ALPHA1, +]; + +/// 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; + } + }; + + match activator + .activate_registered_supervisor(request.clone()) + .await + { + Ok(()) => { + info!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + "Activated warm-pool supervisor from SandboxClaim" + ); + } + Err(status) if status.code() == Code::AlreadyExists => { + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + "Warm-pool supervisor was already activated" + ); + } + Err(status) if status.code() == Code::NotFound => { + debug!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + "Warm-pool supervisor registration is not pending yet" + ); + } + Err(status) => { + warn!( + namespace = %self.namespace, + sandbox_claim = %claim.name, + sandbox = %sandbox_name, + sandbox_id = %request.sandbox_id, + instance_id = %request.instance_id, + code = ?status.code(), + message = %status.message(), + "Failed to activate warm-pool supervisor from SandboxClaim" + ); + } + } + } + + 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 { + for version in SANDBOX_CLAIM_VERSIONS { + let claim_api = self.sandbox_claim_api(client.clone(), version); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + claim_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + Ok(Ok(_)) => return Ok(claim_api), + 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() + )); + } + } + } + Err(format!( + "no supported SandboxClaim API version is available; tried {}", + SANDBOX_CLAIM_VERSIONS.join(", ") + )) + } + + 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); + 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, + uid: String, + warm_pool_name: Option, + sandbox_template_name: Option, + sandbox_name: Option, + pod_ips: Vec, +} + +fn parse_sandbox_claim(obj: &DynamicObject) -> Result { + Ok(ParsedSandboxClaim { + name: required_metadata_field(&obj.metadata.name, "name")?, + uid: required_metadata_field(&obj.metadata.uid, "uid")?, + warm_pool_name: string_at(&obj.data, &["spec", "warmPoolRef", "name"]) + .or_else(|| string_at(&obj.data, &["spec", "warmpool"])), + 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 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}")) +} + +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}" + )); + } + 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 = sandbox_cr + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(LABEL_SANDBOX_ID)) + .filter(|id| !id.is_empty()) + .cloned() + .ok_or_else(|| "Sandbox CR is missing OpenShell sandbox id 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, + 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 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; + + 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.uid = Some(uid.to_string()); + obj.data = data; + obj + } + + 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 + } + + 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.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 parses_v1alpha1_sandbox_claim_status() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1ALPHA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({ + "spec": { + "sandboxTemplateRef": {"name": "template-a"}, + "warmpool": "pool-a" + }, + "status": {"sandbox": {"name": "sandbox-a"}} + }), + ); + + let claim = parse_sandbox_claim(&obj).unwrap(); + + assert_eq!(claim.warm_pool_name.as_deref(), Some("pool-a")); + assert_eq!(claim.sandbox_template_name.as_deref(), Some("template-a")); + assert_eq!(claim.sandbox_name.as_deref(), Some("sandbox-a")); + } + + #[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 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"}} + }), + ); + 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.reason, "SandboxClaim/claim-a"); + } + + #[test] + fn activation_request_rejects_uncontrolled_pod() { + let obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"status": {"sandbox": {"name": "sandbox-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", + "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 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-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index 34d92d84ad..3145af5e7e 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -37,12 +37,15 @@ use openshell_core::proto::{ PlatformEvent, Sandbox, SandboxCondition, SandboxPhase, SandboxSpec, SandboxStatus, SandboxTemplate, ServiceEndpoint, SshSession, }; -use openshell_core::supervisor_bootstrap::SupervisorBootstrapIdentityProvider; +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; @@ -260,6 +263,27 @@ async fn kubernetes_supervisor_bootstrap_identity_provider( } } +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, @@ -398,6 +422,7 @@ pub struct ComputeRuntime { 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, @@ -424,6 +449,7 @@ impl ComputeRuntime { tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, supervisor_bootstrap_identity: Option>, + sandbox_claim_activation: Option>, gateway_bind_addresses: Vec, ) -> Result { let capabilities = driver @@ -457,6 +483,7 @@ impl ComputeRuntime { 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, @@ -519,6 +546,7 @@ impl ComputeRuntime { tracing_log_bus, supervisor_sessions, None, + None, gateway_bind_addresses, ) .await @@ -534,6 +562,8 @@ impl ComputeRuntime { ) -> 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()))?; @@ -550,6 +580,7 @@ impl ComputeRuntime { tracing_log_bus, supervisor_sessions, supervisor_bootstrap_identity, + sandbox_claim_activation, Vec::new(), ) .await @@ -576,6 +607,7 @@ impl ComputeRuntime { tracing_log_bus, supervisor_sessions, None, + None, Vec::new(), ) .await @@ -605,6 +637,7 @@ impl ComputeRuntime { tracing_log_bus, supervisor_sessions, None, + None, Vec::new(), ) .await @@ -627,6 +660,16 @@ impl ComputeRuntime { 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() @@ -2891,6 +2934,7 @@ pub async fn new_test_runtime_for_driver(store: Arc, driver_name: &str) - 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(), @@ -3361,6 +3405,7 @@ mod tests { 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/lib.rs b/crates/openshell-server/src/lib.rs index a8c9f19d95..4a1d379179 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -402,6 +402,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)); From 5aa9e007ed6d068b9a1583d73bf79d80763fbf55 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Thu, 23 Jul 2026 20:28:51 +0100 Subject: [PATCH 07/15] feat(kubernetes): create SandboxClaims for matching warm pools Signed-off-by: Gordon Sim --- Cargo.lock | 1 + architecture/compute-runtimes.md | 29 +- architecture/gateway.md | 13 + crates/openshell-core/src/grpc_client.rs | 74 +- crates/openshell-driver-kubernetes/Cargo.toml | 1 + crates/openshell-driver-kubernetes/README.md | 20 + .../openshell-driver-kubernetes/src/config.rs | 45 + .../openshell-driver-kubernetes/src/driver.rs | 1823 ++++++++++++++++- crates/openshell-driver-kubernetes/src/lib.rs | 2 +- .../openshell-driver-kubernetes/src/main.rs | 3 +- .../src/sandboxclaim.rs | 425 +++- crates/openshell-sandbox/src/lib.rs | 18 +- .../openshell-supervisor-process/src/run.rs | 4 + deploy/helm/openshell/README.md | 1 + .../helm/openshell/templates/clusterrole.yaml | 22 + .../openshell/templates/gateway-config.yaml | 3 + .../templates/warm-pool-profile-role.yaml | 27 + .../warm-pool-profile-rolebinding.yaml | 19 + .../openshell/tests/gateway_config_test.yaml | 9 + .../tests/warm_pool_profiles_rbac_test.yaml | 58 + deploy/helm/openshell/values.yaml | 5 + docs/kubernetes/setup.mdx | 6 +- docs/reference/gateway-config.mdx | 5 + docs/reference/sandbox-compute-drivers.mdx | 10 + 24 files changed, 2456 insertions(+), 167 deletions(-) create mode 100644 deploy/helm/openshell/templates/warm-pool-profile-role.yaml create mode 100644 deploy/helm/openshell/templates/warm-pool-profile-rolebinding.yaml create mode 100644 deploy/helm/openshell/tests/warm_pool_profiles_rbac_test.yaml diff --git a/Cargo.lock b/Cargo.lock index 31e2104987..a7ce66f440 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3944,6 +3944,7 @@ dependencies = [ "prost-types", "serde", "serde_json", + "sha2 0.10.9", "temp-env", "thiserror 2.0.18", "tokio", diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index a072ac0d23..bca1b6b316 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -118,7 +118,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. | @@ -212,6 +212,33 @@ JWT. The supervisor installs that JWT in memory and starts the normal 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 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. +Warm-pool activation uses the same trust boundary as direct Kubernetes +activation. The direct path validates `pod token -> live pod -> owning Sandbox +CR -> Sandbox sandbox-id metadata`; the warm path validates `pod token -> live +pod -> owning Sandbox CR -> associated SandboxClaim -> SandboxClaim sandbox-id +metadata`. In both paths, sandbox-id metadata is not cryptographic proof by +itself. It is trusted because the gateway and Agent Sandbox controllers own the +relevant Kubernetes objects, and RBAC must prevent sandbox workloads and +untrusted users from creating or mutating trusted `Sandbox`, `SandboxClaim`, pod +metadata, or sandbox service-account identity in the gateway-managed namespace. +Under that model, warm pooling does not require a separate anti-spoofing token +or gateway-side claim mapping beyond the same live Kubernetes ownership and +metadata consistency checks used by the direct path. If an operator grants +untrusted principals write access to those objects, both warm and direct +activation require a broader common proof model. +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 defb0019c2..4c14262e12 100644 --- a/architecture/gateway.md +++ b/architecture/gateway.md @@ -169,6 +169,19 @@ 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 sending an activation message with the gateway JWT. +For warm-pooled Kubernetes sandboxes, the same check is re-anchored through the +claim that adopted the warm pod: the live pod must still be controlled by a +live `Sandbox`, that `Sandbox` must be associated with the live +`SandboxClaim`, and the sandbox-id metadata on that claim must identify the +same OpenShell sandbox record before activation can complete. This is not a +separate warm-only spoofing defense; it is the warm-path equivalent of the +direct path's owner-reference and sandbox-id consistency check. Both paths +depend on the same Kubernetes RBAC boundary: sandbox workloads and untrusted +users must not be able to create or mutate trusted Agent Sandbox objects, pod +metadata, or the configured sandbox ServiceAccount in the gateway-managed +namespace. If that boundary is weakened, a fix must cover both direct +`Sandbox` and warm `SandboxClaim` activation rather than adding a warm-only +proof. 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 diff --git a/crates/openshell-core/src/grpc_client.rs b/crates/openshell-core/src/grpc_client.rs index c27ff4dbb7..c261f3eee7 100644 --- a/crates/openshell-core/src/grpc_client.rs +++ b/crates/openshell-core/src/grpc_client.rs @@ -25,7 +25,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use crate::proto::{ DenialSummary, GetDraftPolicyRequest, GetInferenceBundleRequest, GetInferenceBundleResponse, GetSandboxConfigRequest, GetSandboxProviderEnvironmentRequest, NetworkActivitySummary, - PolicyChunk, PolicySource, PolicyStatus, RefreshSandboxTokenRequest, + PodActivationMessage, PolicyChunk, PolicySource, PolicyStatus, RefreshSandboxTokenRequest, RegisterSupervisorPodRequest, ReportPolicyStatusRequest, SandboxPolicy as ProtoSandboxPolicy, SubmitPolicyAnalysisRequest, SubmitPolicyAnalysisResponse, UpdateConfigRequest, inference_client::InferenceClient, open_shell_client::OpenShellClient, @@ -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,6 +292,18 @@ 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}"))? @@ -324,7 +338,61 @@ async fn acquire_k8s_sandbox_token( sandbox_name = %activation.sandbox_name, "received supervisor pod activation" ); - Ok(activation.token) + 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-driver-kubernetes/Cargo.toml b/crates/openshell-driver-kubernetes/Cargo.toml index 2c02f864ab..7ba2794b8a 100644 --- a/crates/openshell-driver-kubernetes/Cargo.toml +++ b/crates/openshell-driver-kubernetes/Cargo.toml @@ -29,6 +29,7 @@ kube-runtime = { workspace = true } k8s-openapi = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +sha2 = { 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 a3ad2166e8..b6a129158a 100644 --- a/crates/openshell-driver-kubernetes/README.md +++ b/crates/openshell-driver-kubernetes/README.md @@ -26,6 +26,26 @@ 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. + ## Workspace Persistence Sandbox pods use a PVC-backed `/sandbox` workspace. An init container seeds the diff --git a/crates/openshell-driver-kubernetes/src/config.rs b/crates/openshell-driver-kubernetes/src/config.rs index 9d13eeb14b..676ca9db58 100644 --- a/crates/openshell-driver-kubernetes/src/config.rs +++ b/crates/openshell-driver-kubernetes/src/config.rs @@ -125,6 +125,20 @@ 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, +} + +impl Default for KubernetesWarmPoolingConfig { + fn default() -> Self { + Self { enabled: true } + } +} + /// Kubernetes `AppArmor` profile requested for the sandbox agent container. #[derive(Debug, Clone, PartialEq, Eq)] pub enum AppArmorProfile { @@ -253,6 +267,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, @@ -346,6 +362,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(), @@ -546,6 +563,34 @@ 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); + } + + #[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_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 8757cb5a8a..79a8c4cae4 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -9,7 +9,7 @@ use crate::config::{ DEFAULT_WORKSPACE_STORAGE_SIZE, KubernetesComputeConfig, SupervisorSideloadMethod, SupervisorTopology, }; -use futures::{Stream, StreamExt, TryStreamExt}; +use futures::{Stream, StreamExt, TryStreamExt, stream}; use k8s_openapi::api::core::v1::{ Event as KubeEventObj, Namespace, Node, PersistentVolumeClaimVolumeSource, Volume, VolumeMount, }; @@ -38,12 +38,13 @@ use openshell_core::proto::compute::v1::{ }; 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 +87,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 +95,16 @@ 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 POD_ANNOTATION_SANDBOX_ID: &str = "openshell.io/sandbox-id"; + const GPU_RESOURCE_NAME: &str = "nvidia.com/gpu"; const SPIFFE_WORKLOAD_API_VOLUME_NAME: &str = "spiffe-workload-api"; @@ -101,6 +113,100 @@ 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, 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 +537,7 @@ pub struct KubernetesComputeDriver { client: Client, watch_client: Client, sandbox_api_version: Arc>, + warm_pool_cache: WarmPoolCache, config: KubernetesComputeConfig, } @@ -476,12 +583,17 @@ 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(); + } + Ok(driver) } pub fn capabilities(&self) -> Result { @@ -525,6 +637,54 @@ 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; + }); + } + 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)) @@ -700,17 +860,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 +925,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,6 +961,381 @@ impl KubernetesComputeDriver { } } + 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() + )); + } + }; + + Ok(list + .items + .into_iter() + .filter_map(|obj| sandbox_from_claim_object(&self.config.namespace, obj).ok()) + .collect()) + } + + 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() + )); + } + }; + + Ok(list + .items + .into_iter() + .find_map(|obj| sandbox_from_claim_object(&self.config.namespace, obj).ok())) + } + + 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() + )); + } + }; + + 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() + )), + } + } + + 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, + all_api.api.list(&ListParams::default().limit(1)), + ) + .await + { + 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)) => 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 { + if !self.config.warm_pooling.enabled { + return Ok(false); + } + + let Some(warm_pool) = self + .matching_warm_pool(target_namespace, request_fingerprint) + .await + else { + return Ok(false); + }; + + 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)) => { + info!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + warm_pool = %warm_pool.name, + "SandboxClaim created in Kubernetes successfully" + ); + Ok(true) + } + 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(false) + } + 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(false) + } + 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(false) + } + } + } + + 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 @@ -862,6 +1404,36 @@ impl KubernetesComputeDriver { 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 self + .try_create_sandbox_claim(sandbox, &target_namespace, &fingerprint) + .await? + { + return Ok(()); + } + 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 @@ -878,7 +1450,7 @@ impl KubernetesComputeDriver { } obj.metadata = ObjectMeta { name: Some(kube_name), - namespace: Some(self.config.namespace.clone()), + namespace: Some(target_namespace), labels: Some(sandbox_labels(sandbox)), annotations: Some(annotations), ..Default::default() @@ -930,6 +1502,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?; @@ -1025,7 +1601,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", @@ -1044,8 +1625,55 @@ 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 mut event_stream = watcher::watcher(event_api, watcher::Config::default()).boxed(); - let (tx, rx) = mpsc::channel(256); + 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); tokio::spawn(async move { let mut sandbox_name_to_id = std::collections::HashMap::::new(); @@ -1108,6 +1736,59 @@ impl KubernetesComputeDriver { break; } }, + result = claim_stream.try_next() => match result { + Ok(Some(Event::Applied(obj))) => { + 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() { + break; + } + } + } + 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( @@ -1149,6 +1830,205 @@ impl KubernetesComputeDriver { } } +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" + ); +} + 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 @@ -1157,55 +2037,396 @@ fn should_try_next_sandbox_api_version(err: &KubeError) -> bool { matches!(err, KubeError::Api(api) if api.code == 404) } -fn validate_gpu_request( - gpu_requirements: Option<&GpuResourceRequirements>, -) -> Result<(), tonic::Status> { - let _ = - effective_driver_gpu_count(gpu_requirements).map_err(tonic::Status::invalid_argument)?; - Ok(()) +fn validate_gpu_request( + gpu_requirements: Option<&GpuResourceRequirements>, +) -> Result<(), tonic::Status> { + let _ = + effective_driver_gpu_count(gpu_requirements).map_err(tonic::Status::invalid_argument)?; + Ok(()) +} + +fn kube_resource_name(workspace: &str, name: &str) -> String { + format!("{workspace}--{name}") +} + +const MAX_KUBE_NAME_LEN: usize = 63; + +fn validate_kube_resource_name_length(workspace: &str, name: &str) -> Result<(), tonic::Status> { + let combined = workspace.len() + 2 + name.len(); // "--" separator + if combined > MAX_KUBE_NAME_LEN { + return Err(tonic::Status::invalid_argument(format!( + "combined Kubernetes resource name '{workspace}--{name}' is {combined} characters, \ + exceeding the DNS-1123 limit of {MAX_KUBE_NAME_LEN}" + ))); + } + Ok(()) +} + +fn sandbox_labels(sandbox: &Sandbox) -> BTreeMap { + let mut labels = BTreeMap::new(); + labels.insert(LABEL_SANDBOX_ID.to_string(), sandbox.id.clone()); + labels.insert(LABEL_SANDBOX_NAME.to_string(), sandbox.name.clone()); + labels.insert( + LABEL_SANDBOX_WORKSPACE.to_string(), + sandbox.workspace.clone(), + ); + labels.insert( + LABEL_MANAGED_BY.to_string(), + LABEL_MANAGED_BY_VALUE.to_string(), + ); + 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_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 kube_resource_name(workspace: &str, name: &str) -> String { - format!("{workspace}--{name}") +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) } -const MAX_KUBE_NAME_LEN: usize = 63; - -fn validate_kube_resource_name_length(workspace: &str, name: &str) -> Result<(), tonic::Status> { - let combined = workspace.len() + 2 + name.len(); // "--" separator - if combined > MAX_KUBE_NAME_LEN { - return Err(tonic::Status::invalid_argument(format!( - "combined Kubernetes resource name '{workspace}--{name}' is {combined} characters, \ - exceeding the DNS-1123 limit of {MAX_KUBE_NAME_LEN}" - ))); +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")); } - Ok(()) + 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 sandbox_labels(sandbox: &Sandbox) -> BTreeMap { - let mut labels = BTreeMap::new(); - labels.insert(LABEL_SANDBOX_ID.to_string(), sandbox.id.clone()); - labels.insert(LABEL_SANDBOX_NAME.to_string(), sandbox.name.clone()); - labels.insert( - LABEL_SANDBOX_WORKSPACE.to_string(), - sandbox.workspace.clone(), - ); - labels.insert( - LABEL_MANAGED_BY.to_string(), - LABEL_MANAGED_BY_VALUE.to_string(), - ); - labels +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 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 extension_api_unavailable(err: &KubeError) -> bool { + matches!(err, KubeError::Api(api) if api.code == 404) } fn sandbox_id_from_object(obj: &DynamicObject) -> Result { @@ -3191,6 +4412,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()? diff --git a/crates/openshell-driver-kubernetes/src/lib.rs b/crates/openshell-driver-kubernetes/src/lib.rs index 8c077fc7e3..bad24c8759 100644 --- a/crates/openshell-driver-kubernetes/src/lib.rs +++ b/crates/openshell-driver-kubernetes/src/lib.rs @@ -13,7 +13,7 @@ pub use bootstrap::{ 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; diff --git a/crates/openshell-driver-kubernetes/src/main.rs b/crates/openshell-driver-kubernetes/src/main.rs index 5e1fc15584..1306b13d4b 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)] @@ -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 index 833dea9bab..2d26d95e6e 100644 --- a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs +++ b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs @@ -22,6 +22,9 @@ 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"; @@ -31,15 +34,11 @@ 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 SANDBOX_CLAIM_VERSION_V1ALPHA1: &str = "v1alpha1"; -const SANDBOX_CLAIM_VERSIONS: &[&str] = &[ - SANDBOX_CLAIM_VERSION_V1BETA1, - SANDBOX_CLAIM_VERSION_V1ALPHA1, -]; /// Watches Kubernetes `SandboxClaim` resources and activates pending /// supervisor bootstrap streams after revalidating live driver state. @@ -217,53 +216,30 @@ impl SandboxClaimActivationController { } }; - match activator - .activate_registered_supervisor(request.clone()) - .await - { - Ok(()) => { - info!( - namespace = %self.namespace, - sandbox_claim = %claim.name, - sandbox = %sandbox_name, - sandbox_id = %request.sandbox_id, - instance_id = %request.instance_id, - "Activated warm-pool supervisor from SandboxClaim" - ); - } - Err(status) if status.code() == Code::AlreadyExists => { - debug!( - namespace = %self.namespace, - sandbox_claim = %claim.name, - sandbox = %sandbox_name, - sandbox_id = %request.sandbox_id, - instance_id = %request.instance_id, - "Warm-pool supervisor was already activated" - ); - } - Err(status) if status.code() == Code::NotFound => { - debug!( - namespace = %self.namespace, - sandbox_claim = %claim.name, - sandbox = %sandbox_name, - sandbox_id = %request.sandbox_id, - instance_id = %request.instance_id, - "Warm-pool supervisor registration is not pending yet" - ); - } - Err(status) => { - warn!( - namespace = %self.namespace, - sandbox_claim = %claim.name, - sandbox = %sandbox_name, - sandbox_id = %request.sandbox_id, - instance_id = %request.instance_id, - code = ?status.code(), - message = %status.message(), - "Failed to activate warm-pool supervisor from SandboxClaim" - ); - } - } + 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 { @@ -274,29 +250,23 @@ impl SandboxClaimActivationController { } async fn supported_sandbox_claim_api(&self, client: Client) -> Result { - for version in SANDBOX_CLAIM_VERSIONS { - let claim_api = self.sandbox_claim_api(client.clone(), version); - match tokio::time::timeout( - KUBE_API_TIMEOUT, - claim_api.api.list(&ListParams::default().limit(1)), - ) - .await - { - Ok(Ok(_)) => return Ok(claim_api), - 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() - )); - } - } + 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() + )), } - Err(format!( - "no supported SandboxClaim API version is available; tried {}", - SANDBOX_CLAIM_VERSIONS.join(", ") - )) } fn sandbox_api(&self, client: Client, version: &'static str) -> Api { @@ -333,6 +303,35 @@ impl SandboxClaimActivationController { }; 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); @@ -374,18 +373,115 @@ struct SandboxClaimApi { struct ParsedSandboxClaim { name: String, uid: String, + sandbox_id: 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")?, uid: required_metadata_field(&obj.metadata.uid, "uid")?, - warm_pool_name: string_at(&obj.data, &["spec", "warmPoolRef", "name"]) - .or_else(|| string_at(&obj.data, &["spec", "warmpool"])), + sandbox_id: obj + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(LABEL_SANDBOX_ID)) + .filter(|id| !id.is_empty()) + .cloned(), + 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"]), @@ -428,7 +524,10 @@ fn activation_request_from_claim_state( .and_then(|labels| labels.get(LABEL_SANDBOX_ID)) .filter(|id| !id.is_empty()) .cloned() - .ok_or_else(|| "Sandbox CR is missing OpenShell sandbox id label".to_string())?; + .or_else(|| claim.sandbox_id.clone()) + .ok_or_else(|| { + "SandboxClaim and selected Sandbox are missing OpenShell sandbox id label".to_string() + })?; let instance_id = pod .metadata .uid @@ -464,6 +563,15 @@ 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 @@ -512,6 +620,44 @@ mod tests { 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, @@ -587,35 +733,12 @@ mod tests { 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 parses_v1alpha1_sandbox_claim_status() { - let obj = dynamic_object( - SANDBOX_CLAIM_GROUP, - SANDBOX_CLAIM_VERSION_V1ALPHA1, - SANDBOX_CLAIM_KIND, - "claim-a", - "claim-uid", - json!({ - "spec": { - "sandboxTemplateRef": {"name": "template-a"}, - "warmpool": "pool-a" - }, - "status": {"sandbox": {"name": "sandbox-a"}} - }), - ); - - let claim = parse_sandbox_claim(&obj).unwrap(); - - assert_eq!(claim.warm_pool_name.as_deref(), Some("pool-a")); - assert_eq!(claim.sandbox_template_name.as_deref(), Some("template-a")); - assert_eq!(claim.sandbox_name.as_deref(), Some("sandbox-a")); - } - #[test] fn claim_without_selected_sandbox_does_not_activate() { let obj = dynamic_object( @@ -684,6 +807,90 @@ mod tests { 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(BTreeMap::from([( + LABEL_SANDBOX_ID.to_string(), + "openshell-sandbox".to_string(), + )])); + 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(), + 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 obj = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({"status": {"sandbox": {"name": "sandbox-a"}}}), + ); + 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 obj = dynamic_object( @@ -726,6 +933,22 @@ mod tests { ); } + #[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 { diff --git a/crates/openshell-sandbox/src/lib.rs b/crates/openshell-sandbox/src/lib.rs index 5ba41b9d7e..3e4cc39afe 100644 --- a/crates/openshell-sandbox/src/lib.rs +++ b/crates/openshell-sandbox/src/lib.rs @@ -105,6 +105,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. @@ -154,6 +156,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, @@ -2076,7 +2091,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-supervisor-process/src/run.rs b/crates/openshell-supervisor-process/src/run.rs index a5ff0456c9..51ea4a70d8 100644 --- a/crates/openshell-supervisor-process/src/run.rs +++ b/crates/openshell-supervisor-process/src/run.rs @@ -108,6 +108,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 d29873c6fc..53a89097be 100644 --- a/deploy/helm/openshell/README.md +++ b/deploy/helm/openshell/README.md @@ -228,6 +228,7 @@ 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}` | Enable transparent Kubernetes warm-pool allocation through Agent Sandbox v1beta1 SandboxClaim resources when a compatible OpenShell-enabled SandboxWarmPool exists in the target 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). | | server.workspaceStorageClass | string | `""` | Kubernetes StorageClass for the workspace PVC in sandbox pods. Empty (default) = omit storageClassName, using the cluster's default StorageClass. Set this on clusters with no default StorageClass, otherwise the workspace PVC stays Pending and the sandbox never starts. | | service.healthPort | int | `8081` | Gateway health service port. | diff --git a/deploy/helm/openshell/templates/clusterrole.yaml b/deploy/helm/openshell/templates/clusterrole.yaml index 073c8835ec..b8401ac850 100644 --- a/deploy/helm/openshell/templates/clusterrole.yaml +++ b/deploy/helm/openshell/templates/clusterrole.yaml @@ -31,3 +31,25 @@ 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 diff --git a/deploy/helm/openshell/templates/gateway-config.yaml b/deploy/helm/openshell/templates/gateway-config.yaml index 41cadfbad4..3344c63226 100644 --- a/deploy/helm/openshell/templates/gateway-config.yaml +++ b/deploy/helm/openshell/templates/gateway-config.yaml @@ -148,6 +148,9 @@ data: supervisor_image_pull_policy = {{ .Values.supervisor.image.pullPolicy | quote }} {{- end }} + [openshell.drivers.kubernetes.warm_pooling] + enabled = {{ .Values.server.warmPooling.enabled }} + [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..215ed1f67a 100644 --- a/deploy/helm/openshell/tests/gateway_config_test.yaml +++ b/deploy/helm/openshell/tests/gateway_config_test.yaml @@ -159,6 +159,15 @@ 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 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 8abd3dea71..8a9f5ae218 100644 --- a/deploy/helm/openshell/values.yaml +++ b/deploy/helm/openshell/values.yaml @@ -198,6 +198,11 @@ 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 # -- 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 27a1d3d45a..6b18c59b01 100644 --- a/docs/reference/gateway-config.mdx +++ b/docs/reference/gateway-config.mdx @@ -298,6 +298,11 @@ 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.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 6700a22c9f..95883b4f88 100644 --- a/docs/reference/sandbox-compute-drivers.mdx +++ b/docs/reference/sandbox-compute-drivers.mdx @@ -345,6 +345,16 @@ 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. + 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. From eee939cd1fbe47566d437c3bf1fe43b182b179c6 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Fri, 24 Jul 2026 19:00:30 +0100 Subject: [PATCH 08/15] feat(kubernetes): add support for automatic warm pool and template creation Signed-off-by: Gordon Sim --- Cargo.lock | 1 + crates/openshell-driver-kubernetes/Cargo.toml | 1 + crates/openshell-driver-kubernetes/README.md | 9 + .../openshell-driver-kubernetes/src/config.rs | 93 +- .../openshell-driver-kubernetes/src/driver.rs | 1292 ++++++++++++++++- deploy/helm/openshell/README.md | 6 +- .../helm/openshell/templates/clusterrole.yaml | 6 + .../openshell/templates/gateway-config.yaml | 6 + .../openshell/tests/gateway_config_test.yaml | 21 + deploy/helm/openshell/values.yaml | 12 + docs/reference/gateway-config.mdx | 7 + docs/reference/sandbox-compute-drivers.mdx | 25 + .../cpu-0-2-configmap.yaml | 18 + .../default-configmap.yaml | 15 + .../env-foo-configmap.yaml | 18 + 15 files changed, 1462 insertions(+), 68 deletions(-) create mode 100644 examples/kubernetes-warm-pool-config/cpu-0-2-configmap.yaml create mode 100644 examples/kubernetes-warm-pool-config/default-configmap.yaml create mode 100644 examples/kubernetes-warm-pool-config/env-foo-configmap.yaml diff --git a/Cargo.lock b/Cargo.lock index a7ce66f440..27fcdb07a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3949,6 +3949,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-stream", + "toml", "tonic", "tracing", "tracing-subscriber", diff --git a/crates/openshell-driver-kubernetes/Cargo.toml b/crates/openshell-driver-kubernetes/Cargo.toml index 7ba2794b8a..6a1a0b2bcc 100644 --- a/crates/openshell-driver-kubernetes/Cargo.toml +++ b/crates/openshell-driver-kubernetes/Cargo.toml @@ -29,6 +29,7 @@ kube-runtime = { workspace = true } k8s-openapi = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +toml = { workspace = true } sha2 = { workspace = true } clap = { workspace = true } tracing = { workspace = true } diff --git a/crates/openshell-driver-kubernetes/README.md b/crates/openshell-driver-kubernetes/README.md index b6a129158a..e4138d0d7f 100644 --- a/crates/openshell-driver-kubernetes/README.md +++ b/crates/openshell-driver-kubernetes/README.md @@ -46,6 +46,15 @@ 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 diff --git a/crates/openshell-driver-kubernetes/src/config.rs b/crates/openshell-driver-kubernetes/src/config.rs index 676ca9db58..2b25d82f37 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")] @@ -131,11 +137,73 @@ 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 } + 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 + } } } @@ -567,6 +635,7 @@ mod tests { fn default_warm_pooling_is_enabled() { let cfg = KubernetesComputeConfig::default(); assert!(cfg.warm_pooling.enabled); + assert!(cfg.warm_pooling.profiles.enabled); } #[test] @@ -580,6 +649,28 @@ mod tests { 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!({ diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index 79a8c4cae4..443372203a 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -11,9 +11,12 @@ use crate::config::{ }; 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,11 +33,11 @@ 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; @@ -103,7 +106,28 @@ 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"; @@ -135,6 +159,56 @@ struct WarmPoolCacheEntry { 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>, @@ -592,6 +666,9 @@ impl KubernetesComputeDriver { }; 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) } @@ -685,6 +762,16 @@ impl KubernetesComputeDriver { }); } + 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)) @@ -751,66 +838,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 { @@ -2029,6 +2062,938 @@ async fn refresh_warm_pool_cache(cache: &WarmPoolCache, client: Client, fallback ); } +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, + workspace_storage_class: &config.workspace_storage_class, + 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() + )), + } +} + +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" + ), + } +} + 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 @@ -7741,6 +8706,93 @@ mod tests { ); } + #[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 workspace_storage_class_omitted_from_cr_spec_when_empty() { let cr = sandbox_to_k8s_spec_for_test( @@ -7753,4 +8805,112 @@ mod tests { .is_none() ); } + + #[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/deploy/helm/openshell/README.md b/deploy/helm/openshell/README.md index 53a89097be..96fa973061 100644 --- a/deploy/helm/openshell/README.md +++ b/deploy/helm/openshell/README.md @@ -228,7 +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}` | 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 | 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). | | server.workspaceStorageClass | string | `""` | Kubernetes StorageClass for the workspace PVC in sandbox pods. Empty (default) = omit storageClassName, using the cluster's default StorageClass. Set this on clusters with no default StorageClass, otherwise the workspace PVC stays Pending and the sandbox never starts. | | service.healthPort | int | `8081` | Gateway health service port. | diff --git a/deploy/helm/openshell/templates/clusterrole.yaml b/deploy/helm/openshell/templates/clusterrole.yaml index b8401ac850..eea70f26ad 100644 --- a/deploy/helm/openshell/templates/clusterrole.yaml +++ b/deploy/helm/openshell/templates/clusterrole.yaml @@ -53,3 +53,9 @@ rules: - 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 3344c63226..e47d3c33a3 100644 --- a/deploy/helm/openshell/templates/gateway-config.yaml +++ b/deploy/helm/openshell/templates/gateway-config.yaml @@ -151,6 +151,12 @@ data: [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/tests/gateway_config_test.yaml b/deploy/helm/openshell/tests/gateway_config_test.yaml index 215ed1f67a..3c4d0e7f30 100644 --- a/deploy/helm/openshell/tests/gateway_config_test.yaml +++ b/deploy/helm/openshell/tests/gateway_config_test.yaml @@ -168,6 +168,27 @@ tests: 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/values.yaml b/deploy/helm/openshell/values.yaml index 8a9f5ae218..5819d70e0d 100644 --- a/deploy/helm/openshell/values.yaml +++ b/deploy/helm/openshell/values.yaml @@ -203,6 +203,18 @@ server: # 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/reference/gateway-config.mdx b/docs/reference/gateway-config.mdx index 6b18c59b01..02c58339c9 100644 --- a/docs/reference/gateway-config.mdx +++ b/docs/reference/gateway-config.mdx @@ -303,6 +303,13 @@ provider_spiffe_workload_api_socket_path = "/spiffe-workload-api/spire-agent.soc # 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 95883b4f88..d29f40fa70 100644 --- a/docs/reference/sandbox-compute-drivers.mdx +++ b/docs/reference/sandbox-compute-drivers.mdx @@ -355,6 +355,31 @@ 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/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" From fdef4d18c28cd0fffa6e5da72f0755c65a96c2c8 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Fri, 24 Jul 2026 22:52:47 +0100 Subject: [PATCH 09/15] test(kubernetes): add warm pool e2e coverage Signed-off-by: Gordon Sim --- e2e/rust/Cargo.toml | 5 + e2e/rust/src/harness/kubernetes.rs | 219 ++++++++ e2e/rust/src/harness/mod.rs | 1 + e2e/rust/tests/kubernetes_warm_pool.rs | 671 +++++++++++++++++++++++++ e2e/with-kube-gateway.sh | 4 + 5 files changed, 900 insertions(+) create mode 100644 e2e/rust/src/harness/kubernetes.rs create mode 100644 e2e/rust/tests/kubernetes_warm_pool.rs diff --git a/e2e/rust/Cargo.toml b/e2e/rust/Cargo.toml index 466a97c573..5cba7c5e4d 100644 --- a/e2e/rust/Cargo.toml +++ b/e2e/rust/Cargo.toml @@ -97,6 +97,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 From 94b29d65573f6a815094d347ab0985e94b3668fb Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Tue, 28 Jul 2026 19:59:42 +0100 Subject: [PATCH 10/15] fix(kubernetes): reconcile ambiguous warm claim creation Signed-off-by: Gordon Sim --- architecture/compute-runtimes.md | 11 + crates/openshell-core/src/error.rs | 5 + .../openshell-driver-kubernetes/src/driver.rs | 252 +++++++++++++++++- .../openshell-driver-kubernetes/src/grpc.rs | 11 + crates/openshell-server/src/compute/mod.rs | 50 ++++ 5 files changed, 321 insertions(+), 8 deletions(-) diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index bca1b6b316..854b09c4ac 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -232,6 +232,17 @@ or gateway-side claim mapping beyond the same live Kubernetes ownership and metadata consistency checks used by the direct path. If an operator grants untrusted principals write access to those objects, both warm and direct activation require a broader common proof model. +Warm claim creation is idempotent by the deterministic claim name. After a +create timeout, transport failure, server error, or conflict, the driver reads +that name back and retries the same claim create when it is not yet visible. It +never switches to direct Sandbox creation after an ambiguous claim write. An +existing claim is accepted only when its OpenShell identity, workspace, +allocation marker, and warm-pool reference match the request. If Kubernetes +still cannot determine whether the claim exists, the driver returns +`Unavailable` and the gateway retains the durable `Provisioning` record. This +prevents another create from assigning a new sandbox ID to the same name while +the original Kubernetes write may still be live; the normal driver watcher can +subsequently reconcile an accepted claim. 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 diff --git a/crates/openshell-core/src/error.rs b/crates/openshell-core/src/error.rs index a149cf006a..1192f97d3c 100644 --- a/crates/openshell-core/src/error.rs +++ b/crates/openshell-core/src/error.rs @@ -119,6 +119,10 @@ pub enum ComputeDriverError { /// A precondition for the operation was not met. #[error("{0}")] Precondition(String), + /// The backend may have accepted the operation, but its result could not be + /// determined. Callers must preserve durable intent for reconciliation. + #[error("{0}")] + Unavailable(String), /// Generic error message. #[error("{0}")] Message(String), @@ -130,6 +134,7 @@ impl From for tonic::Status { ComputeDriverError::AlreadyExists => Self::already_exists("sandbox already exists"), ComputeDriverError::InvalidArgument(m) => Self::invalid_argument(m), ComputeDriverError::Precondition(m) => Self::failed_precondition(m), + ComputeDriverError::Unavailable(m) => Self::unavailable(m), ComputeDriverError::Message(m) => Self::internal(m), } } diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index 443372203a..01170cd546 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -63,6 +63,8 @@ pub enum KubernetesDriverError { #[error("{0}")] Precondition(String), #[error("{0}")] + Unavailable(String), + #[error("{0}")] Message(String), } @@ -81,6 +83,7 @@ impl From for openshell_core::ComputeDriverError { KubernetesDriverError::AlreadyExists => Self::AlreadyExists, KubernetesDriverError::InvalidArgument(m) => Self::InvalidArgument(m), KubernetesDriverError::Precondition(m) => Self::Precondition(m), + KubernetesDriverError::Unavailable(m) => Self::Unavailable(m), KubernetesDriverError::Message(m) => Self::Message(m), } } @@ -90,6 +93,8 @@ 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 CLAIM_CREATE_RECONCILE_ATTEMPTS: usize = 3; +const CLAIM_CREATE_RECONCILE_DELAY: Duration = Duration::from_millis(250); const WARM_POOL_CACHE_RETRY_DELAY: Duration = Duration::from_secs(10); const SANDBOX_GROUP: &str = "agents.x-k8s.io"; @@ -1248,12 +1253,12 @@ impl KubernetesComputeDriver { 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( + let create_result = tokio::time::timeout( KUBE_API_TIMEOUT, claim_api.api.create(&PostParams::default(), &claim), ) - .await - { + .await; + match create_result { Ok(Ok(_result)) => { info!( sandbox_id = %sandbox.id, @@ -1274,8 +1279,17 @@ impl KubernetesComputeDriver { ); Ok(false) } - Ok(Err(err)) if matches!(&err, KubeError::Api(api) if api.code == 409) => { - Err(KubernetesDriverError::from_kube(err)) + Ok(Err(err)) if claim_create_result_is_ambiguous(&err) => { + warn!( + sandbox_id = %sandbox.id, + sandbox_name = %sandbox.name, + namespace = %target_namespace, + warm_pool = %warm_pool.name, + error = %err, + "SandboxClaim create result is ambiguous; reconciling by name" + ); + self.reconcile_sandbox_claim_create(&claim_api.api, &claim) + .await } Ok(Err(err)) => { warn!( @@ -1284,7 +1298,7 @@ impl KubernetesComputeDriver { namespace = %target_namespace, warm_pool = %warm_pool.name, error = %err, - "Failed to create SandboxClaim; falling back to direct Sandbox" + "SandboxClaim create was definitively rejected; falling back to direct Sandbox" ); Ok(false) } @@ -1294,13 +1308,108 @@ impl KubernetesComputeDriver { sandbox_name = %sandbox.name, namespace = %target_namespace, timeout_secs = KUBE_API_TIMEOUT.as_secs(), - "Timed out creating SandboxClaim; falling back to direct Sandbox" + "Timed out creating SandboxClaim; reconciling by name" ); - Ok(false) + self.reconcile_sandbox_claim_create(&claim_api.api, &claim) + .await } } } + async fn reconcile_sandbox_claim_create( + &self, + api: &Api, + desired: &DynamicObject, + ) -> Result { + let name = desired.metadata.name.as_deref().ok_or_else(|| { + KubernetesDriverError::InvalidArgument("SandboxClaim name is required".to_string()) + })?; + + for attempt in 1..=CLAIM_CREATE_RECONCILE_ATTEMPTS { + match tokio::time::timeout(KUBE_API_TIMEOUT, api.get(name)).await { + Ok(Ok(existing)) => { + validate_existing_sandbox_claim(desired, &existing)?; + info!( + sandbox_claim = %name, + attempt, + "Reconciled existing SandboxClaim after ambiguous create" + ); + return Ok(true); + } + Ok(Err(KubeError::Api(err))) if err.code == 404 => { + debug!( + sandbox_claim = %name, + attempt, + "SandboxClaim not visible after ambiguous create; retrying idempotent create" + ); + match tokio::time::timeout( + KUBE_API_TIMEOUT, + api.create(&PostParams::default(), desired), + ) + .await + { + Ok(Ok(_)) => { + info!( + sandbox_claim = %name, + attempt, + "Created SandboxClaim while reconciling ambiguous create" + ); + return Ok(true); + } + Ok(Err(err)) if claim_create_result_is_ambiguous(&err) => { + debug!( + sandbox_claim = %name, + attempt, + error = %err, + "Idempotent SandboxClaim create remains ambiguous" + ); + } + Ok(Err(err)) => { + warn!( + sandbox_claim = %name, + attempt, + error = %err, + "Idempotent SandboxClaim create was rejected after an ambiguous write" + ); + } + Err(_) => { + warn!( + sandbox_claim = %name, + attempt, + timeout_secs = KUBE_API_TIMEOUT.as_secs(), + "Idempotent SandboxClaim create timed out" + ); + } + } + } + Ok(Err(err)) => { + warn!( + sandbox_claim = %name, + attempt, + error = %err, + "Failed to reconcile ambiguous SandboxClaim create" + ); + } + Err(_) => { + warn!( + sandbox_claim = %name, + attempt, + timeout_secs = KUBE_API_TIMEOUT.as_secs(), + "Timed out reconciling ambiguous SandboxClaim create" + ); + } + } + + if attempt < CLAIM_CREATE_RECONCILE_ATTEMPTS { + tokio::time::sleep(CLAIM_CREATE_RECONCILE_DELAY).await; + } + } + + Err(KubernetesDriverError::Unavailable(format!( + "SandboxClaim '{name}' create result remains unknown; preserving provisioning state for reconciliation" + ))) + } + 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}"); @@ -3081,6 +3190,66 @@ fn sandbox_claim_to_k8s_object( obj } +fn claim_create_result_is_ambiguous(err: &KubeError) -> bool { + match err { + KubeError::Api(response) => { + response.code == 408 || response.code == 409 || response.code >= 500 + } + _ => true, + } +} + +fn validate_existing_sandbox_claim( + desired: &DynamicObject, + existing: &DynamicObject, +) -> Result<(), KubernetesDriverError> { + let desired_name = desired.metadata.name.as_deref().unwrap_or_default(); + let existing_name = existing.metadata.name.as_deref().unwrap_or_default(); + if desired_name.is_empty() || existing_name != desired_name { + return Err(KubernetesDriverError::Precondition(format!( + "existing SandboxClaim name '{existing_name}' does not match requested claim '{desired_name}'" + ))); + } + + for key in [ + LABEL_MANAGED_BY, + LABEL_SANDBOX_ID, + LABEL_SANDBOX_NAME, + LABEL_SANDBOX_WORKSPACE, + LABEL_ALLOCATION, + ] { + let expected = desired + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(key)) + .map(String::as_str) + .unwrap_or_default(); + let actual = existing + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(key)) + .map(String::as_str) + .unwrap_or_default(); + if expected.is_empty() || actual != expected { + return Err(KubernetesDriverError::Precondition(format!( + "existing SandboxClaim '{desired_name}' has conflicting {key} metadata" + ))); + } + } + + let expected_pool = string_at(&desired.data, &["spec", "warmPoolRef", "name"]); + let actual_pool = string_at(&existing.data, &["spec", "warmPoolRef", "name"]); + if expected_pool.is_none() || actual_pool != expected_pool { + return Err(KubernetesDriverError::Precondition(format!( + "existing SandboxClaim '{desired_name}' targets a different warm pool" + ))); + } + + Ok(()) +} + fn enabled_warm_pool_from_object(obj: DynamicObject) -> Option { let namespace = obj.metadata.namespace.clone()?; let name = obj.metadata.name.clone()?; @@ -5411,6 +5580,73 @@ mod tests { ); } + #[test] + fn sandbox_claim_create_classifies_conflict_and_server_errors_as_ambiguous() { + assert!(claim_create_result_is_ambiguous(&kube_api_error( + 409, + "already exists" + ))); + assert!(claim_create_result_is_ambiguous(&kube_api_error( + 408, + "request timeout" + ))); + assert!(claim_create_result_is_ambiguous(&kube_api_error( + 500, + "internal error" + ))); + assert!(!claim_create_result_is_ambiguous(&kube_api_error( + 403, + "forbidden" + ))); + assert!(!claim_create_result_is_ambiguous(&kube_api_error( + 422, "invalid" + ))); + } + + #[test] + fn existing_sandbox_claim_is_idempotent_only_for_same_identity_and_pool() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_CLAIM_KIND); + let sandbox = Sandbox { + id: "sandbox-id".to_string(), + name: "dev".to_string(), + workspace: "workspace-a".to_string(), + ..Default::default() + }; + let desired = sandbox_claim_to_k8s_object(&sandbox, "pool-a", &resource); + let mut existing = desired.clone(); + existing.metadata.uid = Some("claim-uid".to_string()); + + validate_existing_sandbox_claim(&desired, &existing) + .expect("same claim identity should be idempotent"); + + existing + .metadata + .labels + .as_mut() + .unwrap() + .insert(LABEL_SANDBOX_ID.to_string(), "other-id".to_string()); + let err = validate_existing_sandbox_claim(&desired, &existing) + .expect_err("different sandbox identity must conflict"); + assert!(matches!(err, KubernetesDriverError::Precondition(_))); + } + + #[test] + fn existing_sandbox_claim_rejects_different_pool() { + let resource = KubernetesComputeDriver::extension_resource(SANDBOX_CLAIM_KIND); + let sandbox = Sandbox { + id: "sandbox-id".to_string(), + name: "dev".to_string(), + workspace: "workspace-a".to_string(), + ..Default::default() + }; + let desired = sandbox_claim_to_k8s_object(&sandbox, "pool-a", &resource); + let existing = sandbox_claim_to_k8s_object(&sandbox, "pool-b", &resource); + + let err = validate_existing_sandbox_claim(&desired, &existing) + .expect_err("different warm pool must conflict"); + assert!(matches!(err, KubernetesDriverError::Precondition(_))); + } + #[test] fn enabled_warm_pool_parse_defaults_template_namespace_to_pool_namespace() { let resource = KubernetesComputeDriver::extension_resource(SANDBOX_WARM_POOL_KIND); diff --git a/crates/openshell-driver-kubernetes/src/grpc.rs b/crates/openshell-driver-kubernetes/src/grpc.rs index fccfa9464b..c7ba35d455 100644 --- a/crates/openshell-driver-kubernetes/src/grpc.rs +++ b/crates/openshell-driver-kubernetes/src/grpc.rs @@ -166,4 +166,15 @@ mod tests { assert_eq!(status.code(), tonic::Code::AlreadyExists); assert_eq!(status.message(), "sandbox already exists"); } + + #[test] + fn ambiguous_driver_errors_map_to_unavailable_status() { + let status: Status = ComputeDriverError::from(KubernetesDriverError::Unavailable( + "create outcome unknown".to_string(), + )) + .into(); + + assert_eq!(status.code(), tonic::Code::Unavailable); + assert_eq!(status.message(), "create outcome unknown"); + } } diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index 3145af5e7e..3c4e1868e3 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -773,6 +773,18 @@ impl ComputeRuntime { self.sandbox_index.remove_sandbox(sandbox.object_id()); Err(Status::failed_precondition(status.message().to_string())) } + Err(status) if status.code() == Code::Unavailable => { + // The driver may have committed the backend create. Keep the + // durable Provisioning row so its watcher can reconcile the + // accepted resource instead of allowing a new sandbox ID to + // reuse the same name. + warn!( + sandbox_id = %sandbox.object_id(), + error = %status, + "sandbox create outcome is ambiguous; preserving provisioning state" + ); + Err(Status::unavailable(status.message().to_string())) + } Err(err) => { let _ = self .store @@ -3074,6 +3086,7 @@ mod tests { struct TestDriver { listed_sandboxes: Vec, current_sandboxes: Vec, + create_error: Option, } #[tonic::async_trait] @@ -3146,6 +3159,9 @@ mod tests { &self, _request: Request, ) -> Result, Status> { + if let Some(code) = self.create_error { + return Err(Status::new(code, "controlled create error")); + } Ok(tonic::Response::new(CreateSandboxResponse {})) } @@ -3442,6 +3458,38 @@ mod tests { sandbox } + #[tokio::test] + async fn ambiguous_create_preserves_provisioning_record() { + let runtime = test_runtime(Arc::new(TestDriver { + create_error: Some(Code::Unavailable), + ..Default::default() + })) + .await; + let sandbox = sandbox_record("sb-ambiguous", "sandbox-a", SandboxPhase::Provisioning); + + let err = runtime + .create_sandbox(sandbox, None) + .await + .expect_err("ambiguous create should remain unavailable"); + + assert_eq!(err.code(), Code::Unavailable); + let stored = runtime + .store + .get_message::("sb-ambiguous") + .await + .unwrap() + .expect("provisioning record must be retained"); + assert_eq!(stored.phase(), SandboxPhase::Provisioning as i32); + assert_eq!(stored.object_name(), "sandbox-a"); + assert_eq!( + runtime + .sandbox_index + .sandbox_id_for_sandbox_name("default", "sandbox-a") + .as_deref(), + Some("sb-ambiguous") + ); + } + fn ssh_session_record(id: &str, sandbox_id: &str) -> SshSession { SshSession { metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { @@ -5725,6 +5773,7 @@ mod tests { }), workspace: "default".to_string(), }], + ..Default::default() })) .await; @@ -5789,6 +5838,7 @@ mod tests { })), workspace: "default".to_string(), }], + ..Default::default() })) .await; From c6e902a0799a7e316ab4d081dce07c1adb4f09f5 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Tue, 28 Jul 2026 20:20:45 +0100 Subject: [PATCH 11/15] fix(kubernetes): reconcile missed warm claim activations Signed-off-by: Gordon Sim --- architecture/compute-runtimes.md | 7 + .../src/sandboxclaim.rs | 360 +++++++++++++++--- 2 files changed, 309 insertions(+), 58 deletions(-) diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index 854b09c4ac..d8458ad959 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -232,6 +232,13 @@ or gateway-side claim mapping beyond the same live Kubernetes ownership and metadata consistency checks used by the direct path. If an operator grants untrusted principals write access to those objects, both warm and direct activation require a broader common proof model. +Claim activation is level-triggered rather than dependent on a single watch +event. The controller periodically relists OpenShell-managed claims and retries +failed watch streams, because a claim can become visible before its selected +`Sandbox`, pod, or outbound supervisor registration. It reconciles different claims +concurrently, deduplicates work by claim UID, and caps concurrency so one late +registration or slow Kubernetes lookup cannot block activation for the +namespace or create unbounded work. Warm claim creation is idempotent by the deterministic claim name. After a create timeout, transport failure, server error, or conflict, the driver reads that name back and retries the same claim create when it is not yet visible. It diff --git a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs index 2d26d95e6e..bfab6af3d9 100644 --- a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs +++ b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs @@ -11,17 +11,22 @@ 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_SANDBOX_ID; +use openshell_core::driver_utils::{LABEL_MANAGED_BY, LABEL_MANAGED_BY_VALUE, LABEL_SANDBOX_ID}; use openshell_core::supervisor_bootstrap::{ SupervisorBootstrapActivationRequest, SupervisorBootstrapActivator, }; +use std::collections::HashSet; use std::sync::Arc; use std::time::Duration; use tokio::sync::watch; +use tokio::task::JoinSet; use tonic::Code; use tracing::{debug, info, warn}; const KUBE_API_TIMEOUT: Duration = Duration::from_secs(30); +const ACTIVATION_RESYNC_INTERVAL: Duration = Duration::from_secs(15); +const ACTIVATION_WATCH_RETRY_DELAY: Duration = Duration::from_secs(5); +const ACTIVATION_MAX_CONCURRENCY: usize = 32; 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); @@ -75,56 +80,155 @@ impl SandboxClaimActivationController { 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; + 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, + retry_after_ms = ACTIVATION_WATCH_RETRY_DELAY.as_millis(), + "SandboxClaim API is not available; warm-pool claim activation will retry" + ); + if wait_for_retry_or_shutdown(ACTIVATION_WATCH_RETRY_DELAY, &mut shutdown_rx) + .await + { + return; } + continue; } - event = stream.try_next() => match event { - Ok(Some(Event::Applied(claim))) => { - self.handle_claim(claim, activator.as_ref()).await; + }; + + let list_api = self + .sandbox_claim_api(self.client.clone(), claim_api.version) + .api; + let managed_selector = format!("{LABEL_MANAGED_BY}={LABEL_MANAGED_BY_VALUE}"); + let watcher_config = watcher::Config::default().labels(&managed_selector); + let mut stream = watcher::watcher(claim_api.api, watcher_config).boxed(); + let mut resync = tokio::time::interval(ACTIVATION_RESYNC_INTERVAL); + resync.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut activation_tasks = JoinSet::new(); + let mut in_flight = HashSet::new(); + let mut completed = HashSet::new(); + + info!( + namespace = %self.namespace, + sandbox_claim_api_version = %claim_api.version, + "Watching Kubernetes SandboxClaims for warm-pool activation" + ); + + let mut restart_watch = false; + while !restart_watch { + tokio::select! { + changed = shutdown_rx.changed() => { + if changed.is_err() || *shutdown_rx.borrow() { + return; + } } - Ok(Some(Event::Restarted(claims))) => { - for claim in claims { - self.handle_claim(claim, activator.as_ref()).await; + event = stream.try_next() => match event { + Ok(Some(Event::Applied(claim))) => { + schedule_claim_activation( + &mut activation_tasks, + &mut in_flight, + &completed, + self.clone(), + activator.clone(), + claim, + ); + } + Ok(Some(Event::Restarted(claims))) => { + for claim in claims { + schedule_claim_activation( + &mut activation_tasks, + &mut in_flight, + &completed, + self.clone(), + activator.clone(), + claim, + ); + } + } + Ok(Some(Event::Deleted(claim))) => { + if let Some(key) = claim_activation_key(&claim) { + completed.remove(&key); + } + } + Ok(None) => { + warn!( + namespace = %self.namespace, + "SandboxClaim watch stream ended; retrying" + ); + restart_watch = true; + } + Err(err) => { + warn!( + namespace = %self.namespace, + error = %err, + "SandboxClaim watch failed; retrying" + ); + restart_watch = true; + } + }, + _ = resync.tick() => { + match tokio::time::timeout( + KUBE_API_TIMEOUT, + list_api.list(&ListParams::default().labels(&managed_selector)), + ) + .await + { + Ok(Ok(claims)) => { + for claim in claims.items { + schedule_claim_activation( + &mut activation_tasks, + &mut in_flight, + &completed, + self.clone(), + activator.clone(), + claim, + ); + } + } + Ok(Err(err)) => warn!( + namespace = %self.namespace, + error = %err, + "Failed to resync SandboxClaims for warm-pool activation" + ), + Err(_) => warn!( + namespace = %self.namespace, + timeout_seconds = KUBE_API_TIMEOUT.as_secs(), + "Timed out resyncing SandboxClaims for warm-pool activation" + ), } } - Ok(Some(Event::Deleted(_))) => {} - Ok(None) => break, - Err(err) => { - warn!( - namespace = %self.namespace, - error = %err, - "SandboxClaim watch failed; warm-pool claim activation stopped" - ); - break; + task_result = activation_tasks.join_next(), if !activation_tasks.is_empty() => { + match task_result { + Some(Ok((key, activated))) => { + in_flight.remove(&key); + if activated { + completed.insert(key); + } + } + Some(Err(err)) => { + warn!( + namespace = %self.namespace, + error = %err, + "SandboxClaim activation task failed; restarting reconciliation" + ); + restart_watch = true; + } + None => {} + } } } } + + activation_tasks.abort_all(); + if wait_for_retry_or_shutdown(ACTIVATION_WATCH_RETRY_DELAY, &mut shutdown_rx).await { + return; + } } } @@ -132,7 +236,7 @@ impl SandboxClaimActivationController { &self, claim: DynamicObject, activator: &(dyn SupervisorBootstrapActivator + '_), - ) { + ) -> bool { let claim = match parse_sandbox_claim(&claim) { Ok(claim) => claim, Err(err) => { @@ -141,7 +245,7 @@ impl SandboxClaimActivationController { error = %err, "Ignoring invalid SandboxClaim" ); - return; + return false; } }; debug!( @@ -160,7 +264,7 @@ impl SandboxClaimActivationController { sandbox_claim = %claim.name, "SandboxClaim has not selected a Sandbox yet" ); - return; + return false; }; let sandbox_cr = match self.get_sandbox_cr(sandbox_name).await { @@ -172,7 +276,7 @@ impl SandboxClaimActivationController { sandbox = %sandbox_name, "SandboxClaim selected Sandbox is not available yet" ); - return; + return false; } Err(err) => { warn!( @@ -182,13 +286,13 @@ impl SandboxClaimActivationController { error = %err, "Failed to read Sandbox selected by SandboxClaim" ); - return; + return false; } }; let pod = match self.resolve_controlled_pod(&sandbox_cr).await { Ok(Some(pod)) => pod, - Ok(None) => return, + Ok(None) => return false, Err(err) => { warn!( namespace = %self.namespace, @@ -197,13 +301,13 @@ impl SandboxClaimActivationController { error = %err, "Failed to resolve Sandbox pod for SandboxClaim activation" ); - return; + return false; } }; let request = match activation_request_from_claim_state(&claim, &sandbox_cr, &pod) { Ok(Some(request)) => request, - Ok(None) => return, + Ok(None) => return false, Err(err) => { warn!( namespace = %self.namespace, @@ -212,7 +316,7 @@ impl SandboxClaimActivationController { error = %err, "SandboxClaim activation validation failed" ); - return; + return false; } }; @@ -239,7 +343,7 @@ impl SandboxClaimActivationController { ACTIVATION_RETRY_INITIAL_DELAY, ACTIVATION_RETRY_MAX_DELAY, ) - .await; + .await } fn sandbox_claim_api(&self, client: Client, version: &'static str) -> SandboxClaimApi { @@ -364,6 +468,56 @@ impl SandboxClaimActivationController { } } +fn schedule_claim_activation( + tasks: &mut JoinSet<(String, bool)>, + in_flight: &mut HashSet, + completed: &HashSet, + controller: SandboxClaimActivationController, + activator: Arc, + claim: DynamicObject, +) { + if tasks.len() >= ACTIVATION_MAX_CONCURRENCY { + return; + } + let Some(key) = begin_claim_activation(in_flight, completed, &claim) else { + return; + }; + tasks.spawn(async move { + let activated = controller.handle_claim(claim, activator.as_ref()).await; + (key, activated) + }); +} + +fn begin_claim_activation( + in_flight: &mut HashSet, + completed: &HashSet, + claim: &DynamicObject, +) -> Option { + let key = claim_activation_key(claim)?; + if completed.contains(&key) { + return None; + } + in_flight.insert(key.clone()).then_some(key) +} + +fn claim_activation_key(claim: &DynamicObject) -> Option { + claim + .metadata + .uid + .clone() + .or_else(|| claim.metadata.name.clone()) +} + +async fn wait_for_retry_or_shutdown( + delay: Duration, + shutdown_rx: &mut watch::Receiver, +) -> bool { + tokio::select! { + () = tokio::time::sleep(delay) => false, + changed = shutdown_rx.changed() => changed.is_err() || *shutdown_rx.borrow(), + } +} + struct SandboxClaimApi { api: Api, version: &'static str, @@ -393,7 +547,7 @@ async fn activate_registered_supervisor_with_retry( attempts: usize, initial_delay: Duration, max_delay: Duration, -) { +) -> bool { let attempts = attempts.max(1); let mut delay = initial_delay; @@ -412,7 +566,7 @@ async fn activate_registered_supervisor_with_retry( attempt, "Activated warm-pool supervisor from SandboxClaim" ); - return; + return true; } Err(status) if status.code() == Code::AlreadyExists => { debug!( @@ -424,7 +578,7 @@ async fn activate_registered_supervisor_with_retry( attempt, "Warm-pool supervisor was already activated" ); - return; + return true; } Err(status) if status.code() == Code::NotFound && attempt < attempts => { debug!( @@ -450,7 +604,7 @@ async fn activate_registered_supervisor_with_retry( attempts, "Warm-pool supervisor registration is not pending after retries" ); - return; + return false; } Err(status) => { warn!( @@ -464,10 +618,12 @@ async fn activate_registered_supervisor_with_retry( message = %status.message(), "Failed to activate warm-pool supervisor from SandboxClaim" ); - return; + return false; } } } + + false } fn parse_sandbox_claim(obj: &DynamicObject) -> Result { @@ -675,6 +831,64 @@ mod tests { obj } + #[test] + fn claim_activation_is_deduplicated_by_uid() { + let first = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({}), + ); + let mut relisted = first.clone(); + relisted.metadata.name = Some("renamed-in-test".to_string()); + let mut in_flight = HashSet::new(); + let mut completed = HashSet::new(); + + assert_eq!( + begin_claim_activation(&mut in_flight, &completed, &first).as_deref(), + Some("claim-uid") + ); + assert_eq!( + begin_claim_activation(&mut in_flight, &completed, &relisted), + None + ); + + in_flight.remove("claim-uid"); + assert_eq!( + begin_claim_activation(&mut in_flight, &completed, &relisted).as_deref(), + Some("claim-uid") + ); + + in_flight.remove("claim-uid"); + completed.insert("claim-uid".to_string()); + assert_eq!( + begin_claim_activation(&mut in_flight, &completed, &relisted), + None + ); + } + + #[test] + fn claim_activation_key_falls_back_to_name() { + let mut claim = dynamic_object( + SANDBOX_CLAIM_GROUP, + SANDBOX_CLAIM_VERSION_V1BETA1, + SANDBOX_CLAIM_KIND, + "claim-a", + "claim-uid", + json!({}), + ); + claim.metadata.uid = None; + let mut in_flight = HashSet::new(); + let completed = HashSet::new(); + + assert_eq!( + begin_claim_activation(&mut in_flight, &completed, &claim).as_deref(), + Some("claim-a") + ); + } + fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str, version: &'static str) -> DynamicObject { let mut obj = dynamic_object( SANDBOX_GROUP, @@ -849,7 +1063,7 @@ mod tests { reason: "SandboxClaim/claim-a".to_string(), }; - activate_registered_supervisor_with_retry( + let activated = activate_registered_supervisor_with_retry( &activator, request, ActivationLogContext { @@ -863,9 +1077,39 @@ mod tests { ) .await; + assert!(activated); assert_eq!(activator.request_count(), 2); } + #[tokio::test] + async fn incomplete_activation_remains_eligible_for_resync() { + let activator = FakeActivator::new(vec![Err(Status::not_found("not pending"))]); + 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(), + reason: "SandboxClaim/claim-a".to_string(), + }; + + let activated = activate_registered_supervisor_with_retry( + &activator, + request, + ActivationLogContext { + namespace: "openshell", + sandbox_claim: "claim-a", + sandbox: "sandbox-a", + }, + 1, + Duration::ZERO, + Duration::ZERO, + ) + .await; + + assert!(!activated); + assert_eq!(activator.request_count(), 1); + } + #[test] fn activation_request_rejects_missing_sandbox_id_labels() { let obj = dynamic_object( From de7aba7fa2550e69945c15035559320a287ad0d9 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Tue, 28 Jul 2026 20:42:44 +0100 Subject: [PATCH 12/15] fix(server): expire activated pod registration tombstones Signed-off-by: Gordon Sim --- architecture/compute-runtimes.md | 4 ++ .../src/supervisor_pod_registration.rs | 53 +++++++++++++++---- 2 files changed, 48 insertions(+), 9 deletions(-) diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index d8458ad959..b6ebf5a2ea 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -239,6 +239,10 @@ failed watch streams, because a claim can become visible before its selected concurrently, deduplicates work by claim UID, and caps concurrency so one late registration or slow Kubernetes lookup cannot block activation for the namespace or create unbounded work. +The gateway retains activated pod-UID tombstones for one hour to reject +duplicate registration and activation races, then prunes them opportunistically +on registry access so long-running gateways do not accumulate one entry per +historical pod. Warm claim creation is idempotent by the deterministic claim name. After a create timeout, transport failure, server error, or conflict, the driver reads that name back and retries the same claim create when it is not yet visible. It diff --git a/crates/openshell-server/src/supervisor_pod_registration.rs b/crates/openshell-server/src/supervisor_pod_registration.rs index 805cfc27a2..d1bafc2513 100644 --- a/crates/openshell-server/src/supervisor_pod_registration.rs +++ b/crates/openshell-server/src/supervisor_pod_registration.rs @@ -16,13 +16,15 @@ 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 std::time::{Duration, Instant}; use tokio::sync::mpsc; use tokio_stream::Stream; use tokio_stream::wrappers::ReceiverStream; use tonic::Status; use tracing::{debug, info, warn}; +const ACTIVATED_TOMBSTONE_TTL: Duration = Duration::from_secs(60 * 60); + #[derive(Debug, Default)] pub struct SupervisorPodRegistrationRegistry { inner: Mutex, @@ -66,9 +68,11 @@ impl SupervisorPodRegistrationRegistry { let instance_name = identity.instance_name.clone(); let owner_name = identity.owner_name.clone(); let driver = identity.driver.clone(); + let now = Instant::now(); let replaced = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + prune_activated_tombstones(&mut inner, now); if inner.activated_instance_id.contains_key(&instance_id) { return Err(Status::already_exists( "supervisor instance has already been activated", @@ -82,7 +86,7 @@ impl SupervisorPodRegistrationRegistry { identity, sender, session_id, - registered_at: Instant::now(), + registered_at: now, }, ) .is_some() @@ -119,7 +123,8 @@ impl SupervisorPodRegistrationRegistry { &self, instance_id: &str, ) -> Result { - let inner = self.inner.lock().expect("pending pod registry poisoned"); + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + prune_activated_tombstones(&mut inner, Instant::now()); if inner.activated_instance_id.contains_key(instance_id) { return Err(Status::already_exists( "supervisor instance has already been activated", @@ -138,8 +143,10 @@ impl SupervisorPodRegistrationRegistry { instance_id: &str, activation: PodActivationMessage, ) -> Result<(), Status> { + let now = Instant::now(); let pending = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + prune_activated_tombstones(&mut inner, now); if inner.activated_instance_id.contains_key(instance_id) { return Err(Status::already_exists( "supervisor instance has already been activated", @@ -156,7 +163,7 @@ impl SupervisorPodRegistrationRegistry { }; inner .activated_instance_id - .insert(instance_id.to_string(), Instant::now()); + .insert(instance_id.to_string(), now); pending }; @@ -187,6 +194,7 @@ impl SupervisorPodRegistrationRegistry { 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"); + prune_activated_tombstones(&mut inner, Instant::now()); if inner.activated_instance_id.contains_key(instance_id) { return Err(Status::already_exists( "supervisor instance has already been activated", @@ -233,11 +241,9 @@ impl SupervisorPodRegistrationRegistry { #[must_use] #[allow(dead_code)] pub fn activated_count(&self) -> usize { - self.inner - .lock() - .expect("pending pod registry poisoned") - .activated_instance_id - .len() + let mut inner = self.inner.lock().expect("pending pod registry poisoned"); + prune_activated_tombstones(&mut inner, Instant::now()); + inner.activated_instance_id.len() } fn remove_if_session(&self, instance_id: &str, session_id: u64) { @@ -265,6 +271,12 @@ impl SupervisorPodRegistrationRegistry { } } +fn prune_activated_tombstones(inner: &mut Inner, now: Instant) { + inner.activated_instance_id.retain(|_, activated_at| { + now.saturating_duration_since(*activated_at) < ACTIVATED_TOMBSTONE_TTL + }); +} + pub struct PendingRegistrationStream { registry: Arc, instance_id: String, @@ -414,6 +426,29 @@ mod tests { assert_eq!(err.code(), tonic::Code::AlreadyExists); } + #[test] + fn expired_activation_tombstones_are_pruned() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + { + let mut inner = registry.inner.lock().expect("registry mutex poisoned"); + inner.activated_instance_id.insert( + "expired-pod-uid".to_string(), + Instant::now() + .checked_sub(ACTIVATED_TOMBSTONE_TTL + Duration::from_secs(1)) + .expect("test tombstone timestamp"), + ); + inner + .activated_instance_id + .insert("fresh-pod-uid".to_string(), Instant::now()); + } + + let stream = registry + .register_pending(bootstrap_identity("expired-pod-uid")) + .expect("expired tombstone must not reject registration"); + assert_eq!(registry.activated_count(), 1); + drop(stream); + } + #[test] fn closed_pending_stream_cannot_be_activated() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); From dab0d887895063d8d4693ba31b55e4e56875ad24 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Tue, 28 Jul 2026 21:27:46 +0100 Subject: [PATCH 13/15] fix(kubernetes): reactivate restarted warm pods Signed-off-by: Gordon Sim --- architecture/compute-runtimes.md | 16 +- .../src/sandboxclaim.rs | 129 ++++++++--- crates/openshell-server/src/compute/mod.rs | 3 +- crates/openshell-server/src/lib.rs | 2 + .../src/supervisor_pod_registration.rs | 205 +++++++++++++----- .../src/warm_pod_activation.rs | 179 +++++++++++++-- 6 files changed, 426 insertions(+), 108 deletions(-) diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index b6ebf5a2ea..70d1281ab2 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -235,14 +235,20 @@ activation require a broader common proof model. Claim activation is level-triggered rather than dependent on a single watch event. The controller periodically relists OpenShell-managed claims and retries failed watch streams, because a claim can become visible before its selected -`Sandbox`, pod, or outbound supervisor registration. It reconciles different claims -concurrently, deduplicates work by claim UID, and caps concurrency so one late -registration or slow Kubernetes lookup cannot block activation for the -namespace or create unbounded work. +`Sandbox`, pod, or outbound supervisor registration. A newly authenticated +warm-pod registration clears completed-claim hints and triggers an immediate +relist, so both a restarted process with the same pod UID and a replacement pod +with a new UID can reactivate against an existing claim. It reconciles different +claims concurrently, deduplicates work by claim UID, and caps concurrency so +one late registration or slow Kubernetes lookup cannot block activation for +the namespace or create unbounded work. The gateway retains activated pod-UID tombstones for one hour to reject duplicate registration and activation races, then prunes them opportunistically on registry access so long-running gateways do not accumulate one entry per -historical pod. +historical pod. A new authenticated registration supersedes its pod-UID +tombstone and receives a new local session ID. Activation and failure delivery +are conditional on that session ID, preventing work started for an old process +from consuming or terminating its replacement stream. Warm claim creation is idempotent by the deterministic claim name. After a create timeout, transport failure, server error, or conflict, the driver reads that name back and retries the same claim create when it is not yet visible. It diff --git a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs index bfab6af3d9..00b2048491 100644 --- a/crates/openshell-driver-kubernetes/src/sandboxclaim.rs +++ b/crates/openshell-driver-kubernetes/src/sandboxclaim.rs @@ -67,17 +67,21 @@ impl SandboxClaimActivationController { pub fn spawn( &self, activator: Arc, + registration_rx: watch::Receiver, shutdown_rx: watch::Receiver, ) { let controller = self.clone(); tokio::spawn(async move { - controller.run(activator, shutdown_rx).await; + controller + .run(activator, registration_rx, shutdown_rx) + .await; }); } async fn run( self, activator: Arc, + mut registration_rx: watch::Receiver, mut shutdown_rx: watch::Receiver, ) { loop { @@ -113,6 +117,7 @@ impl SandboxClaimActivationController { let mut activation_tasks = JoinSet::new(); let mut in_flight = HashSet::new(); let mut completed = HashSet::new(); + let mut registration_notifications_open = true; info!( namespace = %self.namespace, @@ -128,6 +133,27 @@ impl SandboxClaimActivationController { return; } } + changed = registration_rx.changed(), if registration_notifications_open => { + if changed.is_err() { + registration_notifications_open = false; + } else { + invalidate_claim_activation_state( + &mut activation_tasks, + &mut in_flight, + &mut completed, + ); + resync_claim_activations( + &list_api, + &managed_selector, + &mut activation_tasks, + &mut in_flight, + &completed, + self.clone(), + activator.clone(), + ) + .await; + } + } event = stream.try_next() => match event { Ok(Some(Event::Applied(claim))) => { schedule_claim_activation( @@ -173,35 +199,16 @@ impl SandboxClaimActivationController { } }, _ = resync.tick() => { - match tokio::time::timeout( - KUBE_API_TIMEOUT, - list_api.list(&ListParams::default().labels(&managed_selector)), + resync_claim_activations( + &list_api, + &managed_selector, + &mut activation_tasks, + &mut in_flight, + &completed, + self.clone(), + activator.clone(), ) - .await - { - Ok(Ok(claims)) => { - for claim in claims.items { - schedule_claim_activation( - &mut activation_tasks, - &mut in_flight, - &completed, - self.clone(), - activator.clone(), - claim, - ); - } - } - Ok(Err(err)) => warn!( - namespace = %self.namespace, - error = %err, - "Failed to resync SandboxClaims for warm-pool activation" - ), - Err(_) => warn!( - namespace = %self.namespace, - timeout_seconds = KUBE_API_TIMEOUT.as_secs(), - "Timed out resyncing SandboxClaims for warm-pool activation" - ), - } + .await; } task_result = activation_tasks.join_next(), if !activation_tasks.is_empty() => { match task_result { @@ -468,6 +475,57 @@ impl SandboxClaimActivationController { } } +fn invalidate_claim_activation_state( + tasks: &mut JoinSet<(String, bool)>, + in_flight: &mut HashSet, + completed: &mut HashSet, +) { + tasks.abort_all(); + *tasks = JoinSet::new(); + in_flight.clear(); + completed.clear(); +} + +async fn resync_claim_activations( + list_api: &Api, + managed_selector: &str, + tasks: &mut JoinSet<(String, bool)>, + in_flight: &mut HashSet, + completed: &HashSet, + controller: SandboxClaimActivationController, + activator: Arc, +) { + match tokio::time::timeout( + KUBE_API_TIMEOUT, + list_api.list(&ListParams::default().labels(managed_selector)), + ) + .await + { + Ok(Ok(claims)) => { + for claim in claims.items { + schedule_claim_activation( + tasks, + in_flight, + completed, + controller.clone(), + activator.clone(), + claim, + ); + } + } + Ok(Err(err)) => warn!( + namespace = %controller.namespace, + error = %err, + "Failed to resync SandboxClaims for warm-pool activation" + ), + Err(_) => warn!( + namespace = %controller.namespace, + timeout_seconds = KUBE_API_TIMEOUT.as_secs(), + "Timed out resyncing SandboxClaims for warm-pool activation" + ), + } +} + fn schedule_claim_activation( tasks: &mut JoinSet<(String, bool)>, in_flight: &mut HashSet, @@ -889,6 +947,19 @@ mod tests { ); } + #[test] + fn registration_change_invalidates_claim_activation_state() { + let mut tasks = JoinSet::new(); + let mut in_flight = HashSet::from(["claim-in-flight".to_string()]); + let mut completed = HashSet::from(["claim-complete".to_string()]); + + invalidate_claim_activation_state(&mut tasks, &mut in_flight, &mut completed); + + assert!(tasks.is_empty()); + assert!(in_flight.is_empty()); + assert!(completed.is_empty()); + } + fn sandbox_cr(name: &str, uid: &str, sandbox_id: &str, version: &'static str) -> DynamicObject { let mut obj = dynamic_object( SANDBOX_GROUP, diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index 3c4e1868e3..67972a95ee 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -663,10 +663,11 @@ impl ComputeRuntime { pub(crate) fn spawn_sandbox_claim_activation( &self, activator: Arc, + registration_rx: watch::Receiver, shutdown_rx: watch::Receiver, ) { if let Some(controller) = self.sandbox_claim_activation.clone() { - controller.spawn(activator, shutdown_rx); + controller.spawn(activator, registration_rx, shutdown_rx); } } diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 4a1d379179..a458e80f1b 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -402,8 +402,10 @@ pub(crate) async fn run_server( } state.compute.spawn_watchers(shutdown_rx.clone()); + let supervisor_registration_rx = state.supervisor_pod_registrations.subscribe(); state.compute.spawn_sandbox_claim_activation( Arc::new(warm_pod_activation::GatewaySupervisorBootstrapActivator::new(state.clone())), + supervisor_registration_rx, shutdown_rx.clone(), ); ssh_sessions::spawn_session_reaper(store.clone(), Duration::from_secs(3600)); diff --git a/crates/openshell-server/src/supervisor_pod_registration.rs b/crates/openshell-server/src/supervisor_pod_registration.rs index d1bafc2513..ffbfba33db 100644 --- a/crates/openshell-server/src/supervisor_pod_registration.rs +++ b/crates/openshell-server/src/supervisor_pod_registration.rs @@ -17,7 +17,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, watch}; use tokio_stream::Stream; use tokio_stream::wrappers::ReceiverStream; use tonic::Status; @@ -25,10 +25,22 @@ use tracing::{debug, info, warn}; const ACTIVATED_TOMBSTONE_TTL: Duration = Duration::from_secs(60 * 60); -#[derive(Debug, Default)] +#[derive(Debug)] pub struct SupervisorPodRegistrationRegistry { inner: Mutex, next_session_id: AtomicU64, + registration_generation: watch::Sender, +} + +impl Default for SupervisorPodRegistrationRegistry { + fn default() -> Self { + let (registration_generation, _) = watch::channel(0); + Self { + inner: Mutex::new(Inner::default()), + next_session_id: AtomicU64::new(0), + registration_generation, + } + } } #[derive(Debug, Default)] @@ -45,12 +57,23 @@ struct PendingRegistration { registered_at: Instant, } +#[derive(Debug, Clone)] +pub struct PendingRegistrationSnapshot { + pub identity: SupervisorBootstrapIdentity, + pub session_id: u64, +} + impl SupervisorPodRegistrationRegistry { #[must_use] pub fn new() -> Self { Self::default() } + #[must_use] + pub(crate) fn subscribe(&self) -> watch::Receiver { + self.registration_generation.subscribe() + } + #[allow(clippy::result_large_err)] pub fn register_pending( self: &Arc, @@ -73,11 +96,7 @@ impl SupervisorPodRegistrationRegistry { let replaced = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); prune_activated_tombstones(&mut inner, now); - if inner.activated_instance_id.contains_key(&instance_id) { - return Err(Status::already_exists( - "supervisor instance has already been activated", - )); - } + inner.activated_instance_id.remove(&instance_id); inner .pending_by_instance_id .insert( @@ -91,6 +110,8 @@ impl SupervisorPodRegistrationRegistry { ) .is_some() }; + self.registration_generation + .send_modify(|generation| *generation = generation.wrapping_add(1)); if replaced { info!( @@ -122,7 +143,7 @@ impl SupervisorPodRegistrationRegistry { pub(crate) fn pending_identity( &self, instance_id: &str, - ) -> Result { + ) -> Result { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); prune_activated_tombstones(&mut inner, Instant::now()); if inner.activated_instance_id.contains_key(instance_id) { @@ -133,14 +154,18 @@ impl SupervisorPodRegistrationRegistry { inner .pending_by_instance_id .get(instance_id) - .map(|pending| pending.identity.clone()) + .map(|pending| PendingRegistrationSnapshot { + identity: pending.identity.clone(), + session_id: pending.session_id, + }) .ok_or_else(|| Status::not_found("pending supervisor registration not found")) } #[allow(clippy::result_large_err)] - pub(crate) fn activate( + pub(crate) fn activate_if_session( &self, instance_id: &str, + session_id: u64, activation: PodActivationMessage, ) -> Result<(), Status> { let now = Instant::now(); @@ -152,7 +177,7 @@ impl SupervisorPodRegistrationRegistry { "supervisor instance has already been activated", )); } - let Some(pending) = inner.pending_by_instance_id.remove(instance_id) else { + let Some(current) = inner.pending_by_instance_id.get(instance_id) else { debug!( instance_id = %instance_id, "no pending supervisor registration to activate" @@ -161,6 +186,15 @@ impl SupervisorPodRegistrationRegistry { "pending supervisor registration not found", )); }; + if current.session_id != session_id { + return Err(Status::aborted( + "pending supervisor registration was replaced", + )); + } + let pending = inner + .pending_by_instance_id + .remove(instance_id) + .expect("pending registration checked above"); inner .activated_instance_id .insert(instance_id.to_string(), now); @@ -191,7 +225,12 @@ impl SupervisorPodRegistrationRegistry { } #[allow(clippy::result_large_err)] - pub(crate) fn fail_pending(&self, instance_id: &str, status: Status) -> Result<(), Status> { + pub(crate) fn fail_if_session( + &self, + instance_id: &str, + session_id: u64, + status: Status, + ) -> Result<(), Status> { let pending = { let mut inner = self.inner.lock().expect("pending pod registry poisoned"); prune_activated_tombstones(&mut inner, Instant::now()); @@ -200,6 +239,16 @@ impl SupervisorPodRegistrationRegistry { "supervisor instance has already been activated", )); } + let Some(current) = inner.pending_by_instance_id.get(instance_id) else { + return Err(Status::not_found( + "pending supervisor registration not found", + )); + }; + if current.session_id != session_id { + return Err(Status::aborted( + "pending supervisor registration was replaced", + )); + } inner.pending_by_instance_id.remove(instance_id) }; @@ -317,6 +366,16 @@ mod tests { } } + fn activation() -> PodActivationMessage { + 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(), + } + } + #[test] fn dropping_stream_removes_pending_registration() { let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); @@ -356,15 +415,9 @@ mod tests { .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(), - }; + let pending = registry.pending_identity("pod-uid-a").unwrap(); registry - .activate("pod-uid-a", activation) + .activate_if_session("pod-uid-a", pending.session_id, activation()) .expect("activate"); assert_eq!(registry.pending_count(), 0); assert_eq!(registry.activated_count(), 1); @@ -381,49 +434,36 @@ mod tests { #[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) + .activate_if_session("pod-uid-a", 0, 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() { + async fn activation_tombstone_rejects_duplicate_activation_until_reregistration() { 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(), - }; + let pending = registry.pending_identity("pod-uid-a").unwrap(); registry - .activate("pod-uid-a", activation.clone()) + .activate_if_session("pod-uid-a", pending.session_id, activation()) .expect("activate"); let _ = stream.next().await.expect("activation message"); let err = registry - .activate("pod-uid-a", activation) + .activate_if_session("pod-uid-a", pending.session_id, 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); + let replacement = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("new registration supersedes tombstone"); + assert_eq!(registry.activated_count(), 0); + drop(replacement); } #[test] @@ -457,17 +497,80 @@ mod tests { .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) + .activate_if_session("pod-uid-a", 0, activation()) .expect_err("closed stream should remove pending registration"); assert_eq!(err.code(), tonic::Code::NotFound); } + + #[tokio::test] + async fn stale_session_cannot_activate_replacement_stream() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let old_stream = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register old"); + let old = registry.pending_identity("pod-uid-a").unwrap(); + let mut replacement_stream = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register replacement"); + let replacement = registry.pending_identity("pod-uid-a").unwrap(); + + let err = registry + .activate_if_session("pod-uid-a", old.session_id, activation()) + .expect_err("stale session must not activate replacement"); + assert_eq!(err.code(), tonic::Code::Aborted); + assert_eq!(registry.pending_count(), 1); + + registry + .activate_if_session("pod-uid-a", replacement.session_id, activation()) + .expect("activate replacement"); + assert!(replacement_stream.next().await.unwrap().is_ok()); + drop(old_stream); + } + + #[tokio::test] + async fn stale_session_cannot_fail_replacement_stream() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let old_stream = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register old"); + let old = registry.pending_identity("pod-uid-a").unwrap(); + let mut replacement_stream = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register replacement"); + let replacement = registry.pending_identity("pod-uid-a").unwrap(); + + let err = registry + .fail_if_session( + "pod-uid-a", + old.session_id, + Status::permission_denied("stale failure"), + ) + .expect_err("stale session must not fail replacement"); + assert_eq!(err.code(), tonic::Code::Aborted); + + registry + .activate_if_session("pod-uid-a", replacement.session_id, activation()) + .expect("activate replacement"); + assert!(replacement_stream.next().await.unwrap().is_ok()); + drop(old_stream); + } + + #[tokio::test] + async fn registrations_increment_coalescing_generation() { + let registry = Arc::new(SupervisorPodRegistrationRegistry::new()); + let mut generation = registry.subscribe(); + + let first = registry + .register_pending(bootstrap_identity("pod-uid-a")) + .expect("register first"); + let second = registry + .register_pending(bootstrap_identity("pod-uid-b")) + .expect("register second"); + + generation.changed().await.expect("generation change"); + assert_eq!(*generation.borrow_and_update(), 2); + drop((first, second)); + } } diff --git a/crates/openshell-server/src/warm_pod_activation.rs b/crates/openshell-server/src/warm_pod_activation.rs index b93c9307ae..1710fa8a38 100644 --- a/crates/openshell-server/src/warm_pod_activation.rs +++ b/crates/openshell-server/src/warm_pod_activation.rs @@ -17,6 +17,8 @@ use std::sync::Arc; use tonic::Status; use tracing::{info, warn}; +const ACTIVATION_SESSION_RETRY_ATTEMPTS: usize = 3; + #[derive(Debug, Clone, PartialEq, Eq)] #[allow(dead_code)] pub struct WarmPodActivationTarget { @@ -105,30 +107,46 @@ pub async fn activate_warm_pod( 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)?; - mint_pod_activation(state, &validated.sandbox_id, "WarmPodActivation").await - } - .await; - let activation = match activation_result { - Ok(activation) => activation, - Err(status) => { - let stream_status = clone_status(&status); - state - .supervisor_pod_registrations - .fail_pending(&target.instance_id, stream_status)?; - return Err(status); + for _ in 0..ACTIVATION_SESSION_RETRY_ATTEMPTS { + let pending = state + .supervisor_pod_registrations + .pending_identity(&target.instance_id)?; + let activation_result = async { + let validated = validator.validate(&target, &pending.identity).await?; + ensure_validated_activation_matches_pending(&target, &pending.identity, &validated)?; + mint_pod_activation(state, &validated.sandbox_id, "WarmPodActivation").await } - }; + .await; + let activation = match activation_result { + Ok(activation) => activation, + Err(status) => { + let stream_status = clone_status(&status); + match state.supervisor_pod_registrations.fail_if_session( + &target.instance_id, + pending.session_id, + stream_status, + ) { + Ok(()) => return Err(status), + Err(replaced) if replaced.code() == tonic::Code::Aborted => continue, + Err(delivery) => return Err(delivery), + } + } + }; - state - .supervisor_pod_registrations - .activate(&target.instance_id, activation)?; - Ok(()) + match state.supervisor_pod_registrations.activate_if_session( + &target.instance_id, + pending.session_id, + activation, + ) { + Ok(()) => return Ok(()), + Err(replaced) if replaced.code() == tonic::Code::Aborted => {} + Err(status) => return Err(status), + } + } + + Err(Status::aborted( + "supervisor registration changed repeatedly during activation", + )) } fn clone_status(status: &Status) -> Status { @@ -233,6 +251,8 @@ mod tests { use openshell_core::proto::datamodel::v1::ObjectMeta; use openshell_core::proto::{Sandbox, SandboxPhase, SandboxSpec}; use std::collections::HashMap; + use std::sync::Mutex; + use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use tokio_stream::StreamExt; @@ -251,6 +271,32 @@ mod tests { } } + struct ReplaceRegistrationOnceValidator { + state: Arc, + replaced: AtomicBool, + replacement_stream: + Mutex>, + validated: ValidatedWarmPodActivation, + } + + #[async_trait] + impl WarmPodActivationValidator for ReplaceRegistrationOnceValidator { + async fn validate( + &self, + _target: &WarmPodActivationTarget, + _pending: &SupervisorBootstrapIdentity, + ) -> Result { + if !self.replaced.swap(true, Ordering::SeqCst) { + let stream = self + .state + .supervisor_pod_registrations + .register_pending(pending_identity())?; + *self.replacement_stream.lock().unwrap() = Some(stream); + } + Ok(self.validated.clone()) + } + } + async fn state_with_issuer() -> Arc { let mat = generate_jwt_key().expect("jwt key"); let store = Arc::new( @@ -422,6 +468,95 @@ mod tests { assert_eq!(err.code(), tonic::Code::AlreadyExists); } + #[tokio::test] + async fn same_pod_uid_can_reregister_after_activation() { + let state = state_with_issuer().await; + let mut first_stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register first process"); + let validator = StaticValidator { + validated: validated_activation("sandbox-a"), + }; + + activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect("activate first process"); + let first = first_stream.next().await.unwrap().unwrap(); + + let mut restarted_stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("same pod UID must reregister"); + activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect("activate restarted process"); + let restarted = restarted_stream.next().await.unwrap().unwrap(); + + assert_eq!(first.sandbox_id, restarted.sandbox_id); + assert!(!restarted.token.is_empty()); + } + + #[tokio::test] + async fn replacement_pod_uid_can_activate_existing_sandbox() { + let state = state_with_issuer().await; + let mut replacement_identity = pending_identity(); + replacement_identity.instance_id = "pod-uid-b".to_string(); + replacement_identity.instance_name = "warm-pod-b".to_string(); + let mut stream = state + .supervisor_pod_registrations + .register_pending(replacement_identity) + .expect("register replacement pod"); + let validator = StaticValidator { + validated: ValidatedWarmPodActivation { + instance_id: "pod-uid-b".to_string(), + instance_name: "warm-pod-b".to_string(), + ..validated_activation("sandbox-a") + }, + }; + let target = WarmPodActivationTarget { + instance_id: "pod-uid-b".to_string(), + ..activation_target("sandbox-a") + }; + + activate_warm_pod(&state, &validator, target) + .await + .expect("activate replacement pod"); + let activation = stream.next().await.unwrap().unwrap(); + + assert_eq!(activation.sandbox_id, "sandbox-a"); + assert!(!activation.token.is_empty()); + } + + #[tokio::test] + async fn activation_retries_when_registration_changes_during_validation() { + let state = state_with_issuer().await; + let old_stream = state + .supervisor_pod_registrations + .register_pending(pending_identity()) + .expect("register old process"); + let validator = ReplaceRegistrationOnceValidator { + state: state.clone(), + replaced: AtomicBool::new(false), + replacement_stream: Mutex::new(None), + validated: validated_activation("sandbox-a"), + }; + + activate_warm_pod(&state, &validator, activation_target("sandbox-a")) + .await + .expect("replacement session should activate"); + let mut replacement_stream = validator + .replacement_stream + .lock() + .unwrap() + .take() + .expect("replacement stream"); + let activation = replacement_stream.next().await.unwrap().unwrap(); + + assert_eq!(activation.sandbox_id, "sandbox-a"); + drop(old_stream); + } + #[tokio::test] async fn revalidation_metadata_must_match_pending_registration() { let state = state_with_issuer().await; From 6872156e7a5fef95ecee99ab66981a1b10a7e3a9 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Tue, 28 Jul 2026 21:57:09 +0100 Subject: [PATCH 14/15] fix(kubernetes): gate warm pool profiles behind HA lease Signed-off-by: Gordon Sim --- architecture/compute-runtimes.md | 4 + crates/openshell-driver-kubernetes/README.md | 6 + .../openshell-driver-kubernetes/src/driver.rs | 123 ++++++++++++------ crates/openshell-driver-kubernetes/src/lib.rs | 2 +- crates/openshell-server/src/compute/mod.rs | 121 ++++++++++++++++- 5 files changed, 214 insertions(+), 42 deletions(-) diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index 70d1281ab2..24c0d9ec10 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -266,6 +266,10 @@ reconciler that turns admin-authored TOML profiles into generated 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. +The gateway starts the profile reconciler only on the replica holding the +shared reconciler lease and stops it before returning that replica to standby. +The read-side warm-pool cache still runs on every replica so local create +requests do not depend on the lease holder. Kubernetes can run the supervisor in the default combined topology or in a sidecar topology. Combined mode keeps network and process supervision in the diff --git a/crates/openshell-driver-kubernetes/README.md b/crates/openshell-driver-kubernetes/README.md index e4138d0d7f..baeed0ad64 100644 --- a/crates/openshell-driver-kubernetes/README.md +++ b/crates/openshell-driver-kubernetes/README.md @@ -55,6 +55,12 @@ and `gpu_count`. The reconciler generates ordinary `SandboxTemplate` and `SandboxWarmPool` resources labelled `openshell.ai/enabled=true`; the existing matching cache then handles allocation unchanged. +The gateway runs the profile reconciler only while its replica owns the shared +reconciler lease. In HA deployments, lease loss cancels the old reconciler +before that replica returns to standby, so generated resources have one active +gateway writer. The allocation cache remains replica-local and runs on every +gateway because sandbox creation reads it locally. + ## Workspace Persistence Sandbox pods use a PVC-backed `/sandbox` workspace. An init container seeds the diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index 01170cd546..d16545cb5e 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -47,7 +47,7 @@ use std::path::Path; use std::pin::Pin; use std::sync::Arc; use std::time::Duration; -use tokio::sync::{OnceCell, RwLock, mpsc}; +use tokio::sync::{OnceCell, RwLock, mpsc, watch}; use tokio_stream::wrappers::ReceiverStream; use tracing::{debug, info, warn}; @@ -620,6 +620,25 @@ pub struct KubernetesComputeDriver { config: KubernetesComputeConfig, } +#[derive(Clone)] +pub struct WarmPoolProfileReconciler { + client: Client, + watch_client: Client, + config: KubernetesComputeConfig, +} + +impl WarmPoolProfileReconciler { + pub async fn run(&self, cancel: watch::Receiver) { + run_warm_pool_profile_reconciler( + self.client.clone(), + self.watch_client.clone(), + self.config.clone(), + cancel, + ) + .await; + } +} + impl std::fmt::Debug for KubernetesComputeDriver { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("KubernetesComputeDriver") @@ -671,13 +690,20 @@ impl KubernetesComputeDriver { }; 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 warm_pool_profile_reconciler(&self) -> Option { + (self.config.warm_pooling.enabled && self.config.warm_pooling.profiles.enabled).then(|| { + WarmPoolProfileReconciler { + client: self.client.clone(), + watch_client: self.watch_client.clone(), + config: self.config.clone(), + } + }) + } + pub fn capabilities(&self) -> Result { Ok(openshell_core::driver_utils::build_capabilities_response( "kubernetes", @@ -767,16 +793,6 @@ impl KubernetesComputeDriver { }); } - 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)) @@ -2175,6 +2191,7 @@ async fn run_warm_pool_profile_reconciler( client: Client, watch_client: Client, config: KubernetesComputeConfig, + mut cancel: watch::Receiver, ) { let profile_namespace = config .warm_pooling @@ -2188,13 +2205,18 @@ async fn run_warm_pool_profile_reconciler( .to_string(); loop { - reconcile_existing_warm_pool_profiles( - client.clone(), - &config, - &profile_namespace, - &label_selector, - ) - .await; + if *cancel.borrow() { + return; + } + tokio::select! { + () = reconcile_existing_warm_pool_profiles( + client.clone(), + &config, + &profile_namespace, + &label_selector, + ) => {} + _ = cancel.changed() => return, + } let api: Api = Api::namespaced(watch_client.clone(), &profile_namespace); let watcher_config = watcher::Config::default().labels(&label_selector); @@ -2204,15 +2226,14 @@ async fn run_warm_pool_profile_reconciler( 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(Some(event)) => { + tokio::select! { + () = reconcile_warm_pool_profile_event( + client.clone(), + &config, + event, + ) => {} + _ = cancel.changed() => return, } } Ok(None) => { @@ -2233,18 +2254,44 @@ async fn run_warm_pool_profile_reconciler( } } () = tokio::time::sleep(WARM_POOL_PROFILE_RESYNC_DELAY) => { - reconcile_existing_warm_pool_profiles( - client.clone(), - &config, - &profile_namespace, - &label_selector, - ) - .await; + tokio::select! { + () = reconcile_existing_warm_pool_profiles( + client.clone(), + &config, + &profile_namespace, + &label_selector, + ) => {} + _ = cancel.changed() => return, + } } + _ = cancel.changed() => return, } } - tokio::time::sleep(WARM_POOL_PROFILE_REQUEUE_DELAY).await; + tokio::select! { + () = tokio::time::sleep(WARM_POOL_PROFILE_REQUEUE_DELAY) => {} + _ = cancel.changed() => return, + } + } +} + +async fn reconcile_warm_pool_profile_event( + client: Client, + config: &KubernetesComputeConfig, + event: Event, +) { + match event { + Event::Applied(config_map) => { + reconcile_warm_pool_profile_config_map(client, config, config_map).await; + } + Event::Deleted(config_map) => { + garbage_collect_warm_pool_profile(client, config, &config_map).await; + } + Event::Restarted(config_maps) => { + for config_map in config_maps { + reconcile_warm_pool_profile_config_map(client.clone(), config, config_map).await; + } + } } } diff --git a/crates/openshell-driver-kubernetes/src/lib.rs b/crates/openshell-driver-kubernetes/src/lib.rs index bad24c8759..bf536b217d 100644 --- a/crates/openshell-driver-kubernetes/src/lib.rs +++ b/crates/openshell-driver-kubernetes/src/lib.rs @@ -15,6 +15,6 @@ pub use config::{ DEFAULT_WORKSPACE_STORAGE_SIZE, KubernetesComputeConfig, KubernetesSidecarConfig, KubernetesWarmPoolingConfig, SupervisorSideloadMethod, SupervisorTopology, }; -pub use driver::{KubernetesComputeDriver, KubernetesDriverError}; +pub use driver::{KubernetesComputeDriver, KubernetesDriverError, WarmPoolProfileReconciler}; pub use grpc::ComputeDriverService; pub use sandboxclaim::SandboxClaimActivationController; diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index 67972a95ee..56381c4370 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -45,7 +45,7 @@ use openshell_driver_docker::DockerComputeDriver; use openshell_driver_kubernetes::{ ComputeDriverService as KubernetesDriverService, KubernetesComputeDriver, KubernetesSupervisorBootstrapIdentityProvider, LiveK8sResolver, - SandboxClaimActivationController, + SandboxClaimActivationController, WarmPoolProfileReconciler, }; use openshell_driver_podman::{ComputeDriverService as PodmanDriverService, PodmanComputeDriver}; use prost::Message; @@ -194,6 +194,18 @@ trait StartupResume: Send + Sync { async fn resume_sandbox(&self, sandbox_id: &str, sandbox_name: &str) -> Result; } +#[tonic::async_trait] +trait LeaseScopedReconciler: Send + Sync { + async fn run(&self, cancel: watch::Receiver); +} + +#[tonic::async_trait] +impl LeaseScopedReconciler for WarmPoolProfileReconciler { + async fn run(&self, cancel: watch::Receiver) { + self.run(cancel).await; + } +} + #[tonic::async_trait] impl StartupResume for DockerComputeDriver { async fn resume_sandbox(&self, sandbox_id: &str, sandbox_name: &str) -> Result { @@ -423,6 +435,7 @@ pub struct ComputeRuntime { supervisor_sessions: Arc, supervisor_bootstrap_identity: Option>, sandbox_claim_activation: Option>, + lease_scoped_reconciler: Option>, sync_lock: Arc>, delete_gates: Arc, gateway_bind_addresses: Vec, @@ -484,6 +497,7 @@ impl ComputeRuntime { supervisor_sessions, supervisor_bootstrap_identity, sandbox_claim_activation, + lease_scoped_reconciler: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses, @@ -567,8 +581,11 @@ impl ComputeRuntime { let driver = KubernetesComputeDriver::new(config) .await .map_err(|err| ComputeError::Message(err.to_string()))?; + let lease_scoped_reconciler = driver + .warm_pool_profile_reconciler() + .map(|reconciler| -> Arc { Arc::new(reconciler) }); let driver: SharedComputeDriver = Arc::new(KubernetesDriverService::new(driver)); - Self::from_driver( + let mut runtime = Self::from_driver( ComputeDriverKind::Kubernetes.as_str().to_string(), driver, None, @@ -583,7 +600,9 @@ impl ComputeRuntime { sandbox_claim_activation, Vec::new(), ) - .await + .await?; + runtime.lease_scoped_reconciler = lease_scoped_reconciler; + Ok(runtime) } pub(crate) async fn new_remote_driver( @@ -1356,6 +1375,7 @@ impl ComputeRuntime { pub fn spawn_watchers(&self, shutdown_rx: watch::Receiver) { let runtime = Arc::new(self.clone()); if self.store.is_single_replica() { + let _lease_scoped_handle = self.spawn_lease_scoped_reconciler(shutdown_rx.clone()); let watch_runtime = runtime.clone(); let watch_shutdown = shutdown_rx.clone(); tokio::spawn(async move { @@ -1371,6 +1391,17 @@ impl ComputeRuntime { } } + fn spawn_lease_scoped_reconciler( + &self, + cancel: watch::Receiver, + ) -> Option> { + self.lease_scoped_reconciler.clone().map(|reconciler| { + tokio::spawn(async move { + reconciler.run(cancel).await; + }) + }) + } + pub async fn cleanup_on_shutdown(&self) -> Result<(), String> { let Some(cleanup) = &self.shutdown_cleanup else { return Ok(()); @@ -1575,6 +1606,8 @@ impl ComputeRuntime { runtime.reconcile_loop(cancel_rx).await; }); + let lease_scoped_handle = self.spawn_lease_scoped_reconciler(cancel_tx.subscribe()); + loop { tokio::select! { () = tokio::time::sleep(LEASE_RENEWAL_INTERVAL) => { @@ -1601,6 +1634,9 @@ impl ComputeRuntime { let _ = cancel_tx.send(true); let _ = watch_handle.await; let _ = reconcile_handle.await; + if let Some(handle) = lease_scoped_handle { + let _ = handle.await; + } return; } } @@ -1610,6 +1646,9 @@ impl ComputeRuntime { let _ = cancel_tx.send(true); let _ = watch_handle.await; let _ = reconcile_handle.await; + if let Some(handle) = lease_scoped_handle { + let _ = handle.await; + } info!(replica = %lease.replica_id(), "reconciler lease lost — returning to standby"); } @@ -2948,6 +2987,7 @@ pub async fn new_test_runtime_for_driver(store: Arc, driver_name: &str) - supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), supervisor_bootstrap_identity: None, sandbox_claim_activation: None, + lease_scoped_reconciler: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses: Vec::new(), @@ -2970,6 +3010,22 @@ mod tests { use tokio::sync::{Notify, Semaphore, mpsc, oneshot}; use tokio_stream::wrappers::UnboundedReceiverStream; + struct TestLeaseScopedReconciler { + started: Arc, + stopped: Arc, + } + + #[tonic::async_trait] + impl LeaseScopedReconciler for TestLeaseScopedReconciler { + async fn run(&self, mut cancel: watch::Receiver) { + self.started.notify_one(); + if !*cancel.borrow() { + let _ = cancel.changed().await; + } + self.stopped.notify_one(); + } + } + fn string_value(value: &str) -> prost_types::Value { prost_types::Value { kind: Some(prost_types::value::Kind::StringValue(value.to_string())), @@ -3423,6 +3479,7 @@ mod tests { supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), supervisor_bootstrap_identity: None, sandbox_claim_activation: None, + lease_scoped_reconciler: None, sync_lock: Arc::new(Mutex::new(())), delete_gates: Arc::new(DeleteGateRegistry::default()), gateway_bind_addresses: Vec::new(), @@ -3430,6 +3487,64 @@ mod tests { } } + #[tokio::test] + async fn lease_scoped_reconciler_stops_when_cancelled() { + let mut runtime = test_runtime(Arc::new(TestDriver::default())).await; + let started = Arc::new(Notify::new()); + let stopped = Arc::new(Notify::new()); + runtime.lease_scoped_reconciler = Some(Arc::new(TestLeaseScopedReconciler { + started: started.clone(), + stopped: stopped.clone(), + })); + + let (cancel_tx, cancel_rx) = watch::channel(false); + let handle = runtime + .spawn_lease_scoped_reconciler(cancel_rx) + .expect("configured reconciler should start"); + tokio::time::timeout(Duration::from_secs(1), started.notified()) + .await + .expect("reconciler should start"); + + cancel_tx.send(true).unwrap(); + handle.await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), stopped.notified()) + .await + .expect("reconciler should stop after cancellation"); + } + + #[tokio::test] + async fn lease_holder_runs_and_awaits_lease_scoped_reconciler() { + let driver = ControlledDriver::new(); + let mut runtime = test_runtime(driver).await; + let started = Arc::new(Notify::new()); + let stopped = Arc::new(Notify::new()); + runtime.lease_scoped_reconciler = Some(Arc::new(TestLeaseScopedReconciler { + started: started.clone(), + stopped: stopped.clone(), + })); + + let lease = lease::ReconcilerLease::new( + runtime.store.clone(), + runtime.replica_id.clone(), + lease::LEASE_TTL, + ); + let guard = lease.acquire_or_steal().await.unwrap(); + let (shutdown_tx, mut shutdown_rx) = watch::channel(false); + let runtime = Arc::new(runtime); + let holder = tokio::spawn(async move { + runtime.run_as_holder(&lease, guard, &mut shutdown_rx).await; + }); + + tokio::time::timeout(Duration::from_secs(1), started.notified()) + .await + .expect("lease holder should start the reconciler"); + shutdown_tx.send(true).unwrap(); + holder.await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), stopped.notified()) + .await + .expect("lease holder should await reconciler shutdown"); + } + fn register_test_supervisor_session(runtime: &ComputeRuntime, sandbox_id: &str) { let (tx, _rx) = mpsc::channel(1); let (shutdown_tx, _shutdown_rx) = oneshot::channel(); From 7c11c99cdf20a35dc6b3d6227958460b11fcddc1 Mon Sep 17 00:00:00 2001 From: Gordon Sim Date: Wed, 29 Jul 2026 08:42:25 +0100 Subject: [PATCH 15/15] refactor(kubernetes): centralize workspace namespace mapping Signed-off-by: Gordon Sim --- architecture/compute-runtimes.md | 5 +- crates/openshell-driver-kubernetes/README.md | 11 +- .../openshell-driver-kubernetes/src/driver.rs | 232 +++++++++++++++--- docs/reference/sandbox-compute-drivers.mdx | 5 +- 4 files changed, 211 insertions(+), 42 deletions(-) diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index 24c0d9ec10..1a86fa6503 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -265,7 +265,10 @@ 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. +resulting enabled warm-pool resources and their fingerprints. Sandbox creation +and profile reconciliation share one workspace-to-namespace mapping boundary; +the current shared-namespace implementation maps every workspace to the +configured Kubernetes namespace. The gateway starts the profile reconciler only on the replica holding the shared reconciler lease and stops it before returning that replica to standby. The read-side warm-pool cache still runs on every replica so local create diff --git a/crates/openshell-driver-kubernetes/README.md b/crates/openshell-driver-kubernetes/README.md index baeed0ad64..8ce0443e88 100644 --- a/crates/openshell-driver-kubernetes/README.md +++ b/crates/openshell-driver-kubernetes/README.md @@ -34,9 +34,11 @@ allocation. When `warm_pooling.enabled` is true, the driver watches and caches `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. +On create, the driver maps the OpenShell workspace to its target Kubernetes +namespace, renders the spec that direct `Sandbox` creation would use, strips +per-sandbox identity values from the fingerprint, and matches it against the +warm-pool cache for that namespace. The current shared-namespace mapping sends +every workspace to the configured 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 @@ -51,7 +53,8 @@ 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 +and `gpu_count`. The reconciler uses the same workspace-to-namespace mapping as +sandbox creation and generates ordinary `SandboxTemplate` and `SandboxWarmPool` resources labelled `openshell.ai/enabled=true`; the existing matching cache then handles allocation unchanged. diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index d16545cb5e..e9c2627a03 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -225,6 +225,17 @@ struct WarmPoolCacheState { entries: BTreeMap<(String, String), WarmPoolCacheEntry>, } +/// Map an `OpenShell` workspace to the Kubernetes namespace that owns its +/// compute resources. +/// +/// `OpenShell` currently operates the Kubernetes driver in shared-namespace +/// mode, so every workspace resolves to the configured namespace. Keeping the +/// mapping behind this function gives future managed and operator workspace +/// modes one boundary to replace. +fn workspace_to_namespace(config: &KubernetesComputeConfig, _workspace: &str) -> String { + config.namespace.clone() +} + #[derive(Debug, Eq, PartialEq)] enum WarmPoolCacheLookup { NotReady, @@ -738,10 +749,14 @@ impl KubernetesComputeDriver { ) } - fn agent_sandbox_api(&self, client: Client, sandbox_api_version: &str) -> AgentSandboxApi { + fn agent_sandbox_api( + client: Client, + namespace: &str, + sandbox_api_version: &str, + ) -> AgentSandboxApi { let gvk = GroupVersionKind::gvk(SANDBOX_GROUP, sandbox_api_version, SANDBOX_KIND); let resource = ApiResource::from_gvk(&gvk); - let api = Api::namespaced_with(client, &self.config.namespace, &resource); + let api = Api::namespaced_with(client, namespace, &resource); AgentSandboxApi { api, resource } } @@ -794,15 +809,35 @@ impl KubernetesComputeDriver { } 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)) + self.supported_agent_sandbox_api_for_namespace(client, &self.config.namespace) + .await + } + + async fn supported_agent_sandbox_api_for_namespace( + &self, + client: Client, + namespace: &str, + ) -> Result { + let sandbox_api_version = self + .supported_sandbox_api_version(client.clone(), namespace) + .await?; + Ok(Self::agent_sandbox_api( + client, + namespace, + sandbox_api_version, + )) } - async fn supported_sandbox_api_version(&self, client: Client) -> Result<&'static str, String> { + async fn supported_sandbox_api_version( + &self, + client: Client, + namespace: &str, + ) -> Result<&'static str, String> { self.sandbox_api_version - .get_or_try_init( - || async move { self.detect_supported_sandbox_api_version(client).await }, - ) + .get_or_try_init(|| async move { + self.detect_supported_sandbox_api_version(client, namespace) + .await + }) .await .copied() } @@ -810,9 +845,11 @@ impl KubernetesComputeDriver { async fn detect_supported_sandbox_api_version( &self, client: Client, + namespace: &str, ) -> Result<&'static str, String> { for sandbox_api_version in SANDBOX_VERSIONS { - let agent_sandbox_api = self.agent_sandbox_api(client.clone(), sandbox_api_version); + let agent_sandbox_api = + Self::agent_sandbox_api(client.clone(), namespace, sandbox_api_version); match tokio::time::timeout( KUBE_API_TIMEOUT, agent_sandbox_api.api.list(&ListParams::default().limit(1)), @@ -821,7 +858,7 @@ impl KubernetesComputeDriver { { Ok(Ok(_)) => { debug!( - namespace = %self.config.namespace, + namespace, sandbox_api_version = %sandbox_api_version, "Selected Agent Sandbox API version" ); @@ -829,7 +866,7 @@ impl KubernetesComputeDriver { } Ok(Err(err)) if should_try_next_sandbox_api_version(&err) => { debug!( - namespace = %self.config.namespace, + namespace, sandbox_api_version = %sandbox_api_version, error = %err, "Sandbox API version is not available; trying next supported version" @@ -858,13 +895,11 @@ impl KubernetesComputeDriver { /// `openshift.io/sa.scc.uid-range` / `openshift.io/sa.scc.supplemental-groups` /// annotations. /// - If neither config nor `OpenShift` is found, returns `(1000, 1000, {})` as defaults. - async fn resolve_sandbox_identity(&self) -> (u32, u32, BTreeMap) { - resolve_sandbox_identity_for_config( - self.client.clone(), - &self.config, - self.config.namespace.as_str(), - ) - .await + async fn resolve_sandbox_identity( + &self, + namespace: &str, + ) -> (u32, u32, BTreeMap) { + resolve_sandbox_identity_for_config(self.client.clone(), &self.config, namespace).await } async fn has_gpu_capacity(&self) -> Result { @@ -1509,21 +1544,23 @@ impl KubernetesComputeDriver { .map_err(KubernetesDriverError::InvalidArgument)?; let name = sandbox.name.as_str(); + let target_namespace = workspace_to_namespace(&self.config, &sandbox.workspace); info!( sandbox_id = %sandbox.id, sandbox_name = %name, - namespace = %self.config.namespace, + workspace = %sandbox.workspace, + namespace = %target_namespace, "Creating sandbox in Kubernetes" ); let agent_sandbox_api = self - .supported_agent_sandbox_api(self.client.clone()) + .supported_agent_sandbox_api_for_namespace(self.client.clone(), &target_namespace) .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; + self.resolve_sandbox_identity(&target_namespace).await; let params = SandboxPodParams { default_image: &self.config.default_image, @@ -1562,7 +1599,6 @@ impl KubernetesComputeDriver { 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) => { @@ -2373,21 +2409,32 @@ async fn reconcile_existing_warm_pool_profiles( .iter() .filter_map(|config_map| config_map.metadata.uid.clone()) .collect::>(); + let mut target_namespaces = list + .items + .iter() + .filter_map(|config_map| warm_pool_profile_workspace(config_map).ok()) + .map(|workspace| workspace_to_namespace(config, &workspace)) + .collect::>(); + // Shared-namespace mode must still sweep the configured namespace + // when there are no valid profiles in the current list. + target_namespaces.insert(workspace_to_namespace(config, "default")); 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" - ); + for target_namespace in target_namespaces { + if let Err(err) = garbage_collect_orphaned_warm_pool_profile_resources( + client.clone(), + &target_namespace, + &active_source_uids, + ) + .await + { + warn!( + namespace = %target_namespace, + error = %err, + "Failed to garbage-collect orphaned warm-pool profile resources" + ); + } } } Ok(Err(err)) => { @@ -2411,6 +2458,7 @@ async fn reconcile_warm_pool_profile_config_map( config: &KubernetesComputeConfig, config_map: ConfigMap, ) { + let previous_target_namespace = warm_pool_profile_recorded_target_namespace(&config_map); 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 { @@ -2442,6 +2490,23 @@ async fn reconcile_warm_pool_profile_config_map( "Failed to garbage-collect superseded warm-pool profile resources" ); } + if let Some(previous_target_namespace) = previous_target_namespace + && previous_target_namespace != rendered.target_namespace + && let Err(err) = garbage_collect_warm_pool_profile_resources( + client.clone(), + &previous_target_namespace, + &rendered.source.uid, + None, + ) + .await + { + warn!( + namespace = %previous_target_namespace, + config_map = %rendered.source.name, + error = %err, + "Failed to garbage-collect warm-pool profile resources from previous workspace namespace" + ); + } patch_warm_pool_profile_status( client, &rendered.source.namespace, @@ -2496,7 +2561,7 @@ async fn render_warm_pool_profile( config.warm_pooling.profiles.effective_max_replicas(), )?; - let target_namespace = config.namespace.clone(); + let target_namespace = workspace_to_namespace(config, &profile.workspace); 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; @@ -2591,6 +2656,37 @@ fn parse_warm_pool_profile(data: &str) -> Result { .map_err(|err| format!("failed to parse {WARM_POOL_PROFILE_DATA_KEY}: {err}")) } +fn warm_pool_profile_workspace(config_map: &ConfigMap) -> Result { + 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}'"))?; + Ok(parse_warm_pool_profile(data)?.workspace) +} + +fn warm_pool_profile_recorded_target_namespace(config_map: &ConfigMap) -> Option { + config_map + .metadata + .annotations + .as_ref() + .and_then(|annotations| annotations.get(ANNOTATION_WARM_POOL_TARGET_NAMESPACE)) + .filter(|namespace| !namespace.is_empty()) + .cloned() +} + +fn warm_pool_profile_cleanup_namespace( + config: &KubernetesComputeConfig, + config_map: &ConfigMap, +) -> String { + warm_pool_profile_recorded_target_namespace(config_map).unwrap_or_else(|| { + warm_pool_profile_workspace(config_map).map_or_else( + |_| workspace_to_namespace(config, "default"), + |workspace| workspace_to_namespace(config, &workspace), + ) + }) +} + fn validate_warm_pool_profile(profile: &WarmPoolProfile, max_replicas: u32) -> Result<(), String> { if profile.version != WARM_POOL_PROFILE_VERSION { return Err(format!( @@ -2950,7 +3046,7 @@ async fn garbage_collect_warm_pool_profile( let Some(uid) = cm.metadata.uid.as_deref() else { return; }; - let namespace = config.namespace.clone(); + let namespace = warm_pool_profile_cleanup_namespace(config, cm); if let Err(err) = garbage_collect_warm_pool_profile_resources(client, &namespace, uid, None).await { @@ -9021,6 +9117,70 @@ gpu_count = 1 assert_eq!(profile.resources.gpu_count, Some(1)); } + #[test] + fn workspace_to_namespace_uses_shared_configured_namespace() { + let config = KubernetesComputeConfig { + namespace: "shared-sandboxes".to_string(), + ..Default::default() + }; + + assert_eq!( + workspace_to_namespace(&config, "default"), + "shared-sandboxes" + ); + assert_eq!( + workspace_to_namespace(&config, "research"), + "shared-sandboxes" + ); + } + + #[test] + fn warm_pool_profile_cleanup_prefers_recorded_target_namespace() { + let config = KubernetesComputeConfig { + namespace: "shared-sandboxes".to_string(), + ..Default::default() + }; + let config_map = ConfigMap { + metadata: ObjectMeta { + annotations: Some(BTreeMap::from([( + ANNOTATION_WARM_POOL_TARGET_NAMESPACE.to_string(), + "previous-workspace-namespace".to_string(), + )])), + ..Default::default() + }, + data: Some(BTreeMap::from([( + WARM_POOL_PROFILE_DATA_KEY.to_string(), + "version = 1\nworkspace = \"research\"\n".to_string(), + )])), + ..Default::default() + }; + + assert_eq!( + warm_pool_profile_cleanup_namespace(&config, &config_map), + "previous-workspace-namespace" + ); + } + + #[test] + fn warm_pool_profile_cleanup_maps_profile_workspace_without_status() { + let config = KubernetesComputeConfig { + namespace: "shared-sandboxes".to_string(), + ..Default::default() + }; + let config_map = ConfigMap { + data: Some(BTreeMap::from([( + WARM_POOL_PROFILE_DATA_KEY.to_string(), + "version = 1\nworkspace = \"research\"\n".to_string(), + )])), + ..Default::default() + }; + + assert_eq!( + warm_pool_profile_cleanup_namespace(&config, &config_map), + workspace_to_namespace(&config, "research") + ); + } + #[test] fn warm_pool_profile_rejects_unknown_fields() { let err = parse_warm_pool_profile( diff --git a/docs/reference/sandbox-compute-drivers.mdx b/docs/reference/sandbox-compute-drivers.mdx index d29f40fa70..c0d1e50dff 100644 --- a/docs/reference/sandbox-compute-drivers.mdx +++ b/docs/reference/sandbox-compute-drivers.mdx @@ -351,13 +351,16 @@ create requests by creating `extensions.agents.x-k8s.io/v1beta1` 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 workspace. The current shared-namespace mapping selects the configured +`sandbox_namespace` for every 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. +The reconciler uses the same workspace-to-namespace mapping as sandbox +creation. Each profile ConfigMap must contain `data["warm-pool.toml"]` with a narrow CLI-shaped schema: