feat: add warm-pool provisioning for Kubernetes sandboxes

[mrunalp]

Problem Statement

OpenShell should be able to launch many Kubernetes-backed sandboxes quickly by using warm pools. The Kubernetes backend already provisions sandboxes through the upstream Kubernetes SIG Apps Agent Sandbox project, and that project now exposes extension CRDs for SandboxTemplate, SandboxWarmPool, and SandboxClaim. OpenShell currently cold-starts every Kubernetes sandbox by creating a new agents.x-k8s.io Sandbox resource directly, so it does not benefit from pre-warmed pods or pre-provisioned workspace volumes.

Technical Context

The current OpenShell create path is cold-start only. openshell sandbox create builds a public CreateSandboxRequest, the gateway validates and persists a provisioning Sandbox, the compute runtime converts it to the internal DriverSandbox, and the Kubernetes driver renders a single Agent Sandbox Sandbox CR with the final pod template, labels, annotations, supervisor bootstrap environment, projected service account token volume, and workspace PVC templates.

Upstream Agent Sandbox warm pools are modeled through the extension API group extensions.agents.x-k8s.io: SandboxWarmPool keeps ready sandboxes based on a SandboxTemplate, and SandboxClaim checks out an instance from a pool. The important integration constraint is that upstream SandboxClaim.spec.env and SandboxClaim.spec.volumeClaimTemplates force a cold start, so real warm-start support requires OpenShell to put as much stable configuration as possible into the warm-pool template and solve how OpenShell identity is bound after a warm pod has already started.

Affected Components

Component	Key Files	Role
CLI sandbox creation	crates/openshell-cli/src/main.rs, crates/openshell-cli/src/run.rs	Builds CreateSandboxRequest and currently exposes --driver-config-json for driver-specific create options.
Public gateway API	proto/openshell.proto, crates/openshell-server/src/grpc/sandbox.rs	Validates create requests, fills defaults, creates public sandbox metadata, and calls compute.
Compute runtime	crates/openshell-server/src/compute/mod.rs, proto/compute_driver.proto	Persists sandbox state before driver provisioning and forwards selected driver config to the active compute driver.
Kubernetes driver	crates/openshell-driver-kubernetes/src/driver.rs, crates/openshell-driver-kubernetes/src/config.rs	Detects Agent Sandbox APIs, renders pod specs, creates/deletes/lists/watches Kubernetes resources.
Supervisor bootstrap identity	crates/openshell-core/src/sandbox_env.rs, crates/openshell-server/src/auth/k8s_sa.rs, crates/openshell-supervisor-process/src/run.rs	Binds supervisor relay and token bootstrap to OpenShell sandbox ID/name and pod annotations.
Helm and local k8s dev	deploy/helm/openshell/templates/role.yaml, deploy/helm/openshell/templates/gateway-config.yaml, deploy/helm/openshell/values.yaml, tasks/scripts/helm-k3s-local.sh	Renders driver config, RBAC, and local Agent Sandbox installation.
Docs	docs/reference/gateway-config.mdx, docs/reference/sandbox-compute-drivers.mdx, docs/kubernetes/setup.mdx, crates/openshell-driver-kubernetes/README.md, architecture/sandbox.md	Documents Kubernetes driver configuration, Agent Sandbox setup, and sandbox startup invariants.

Technical Investigation

Architecture Overview

Current create flow:

1. openshell sandbox create parses create flags in crates/openshell-cli/src/main.rs and sends a CreateSandboxRequest from crates/openshell-cli/src/run.rs.
2. handle_create_sandbox_inner in crates/openshell-server/src/grpc/sandbox.rs validates the request, resolves the default image, applies policy validation, creates a gateway sandbox ID/name, and calls state.compute.create_sandbox.
3. ComputeRuntime::create_sandbox in crates/openshell-server/src/compute/mod.rs converts the public sandbox to a driver sandbox, persists the public sandbox record before driver provisioning, optionally injects a sandbox token for local drivers, and calls the active driver.
4. KubernetesDriver::create_sandbox in crates/openshell-driver-kubernetes/src/driver.rs validates Kubernetes driver config and GPU requirements, resolves sandbox UID/GID, renders one agents.x-k8s.io Sandbox dynamic object, and creates it through the Kubernetes API.
5. The Agent Sandbox controller creates the pod. The OpenShell supervisor starts in that pod with OPENSHELL_SANDBOX_ID, OPENSHELL_SANDBOX, OPENSHELL_ENDPOINT, TLS/material mounts, and the projected service account token needed for IssueSandboxToken.

Warm-pool flow would add a second Kubernetes object graph. Instead of creating a final Sandbox directly, OpenShell would create or reference a SandboxTemplate and SandboxWarmPool, then create a SandboxClaim for each OpenShell sandbox. The claim controller would adopt a pre-warmed Sandbox when possible. OpenShell would then need to map the claim, adopted Sandbox, pod, and Kubernetes events back to the persisted OpenShell sandbox ID.

Code References

Location	Description
proto/openshell.proto:443	Public CreateSandboxRequest has spec, optional name, and labels; no warm-pool-specific public field today.
proto/compute_driver.proto:130	Driver-owned DriverSandboxTemplate.driver_config is the established pass-through for selected driver-specific create options.
crates/openshell-cli/src/main.rs:1279	CLI exposes --driver-config-json, currently the narrowest place to pass Kubernetes-specific warm-pool selection without a new public flag.
crates/openshell-cli/src/run.rs:1903	CLI parses driver config JSON and includes it in the sandbox template before sending CreateSandboxRequest.
crates/openshell-server/src/grpc/sandbox.rs:154	Server ensures the sandbox template always carries the resolved image before validation/create.
crates/openshell-server/src/grpc/sandbox.rs:203	Kubernetes sandboxes ignore the gateway-minted token and bootstrap through IssueSandboxToken.
crates/openshell-server/src/compute/mod.rs:480	Compute runtime create entry point persists the sandbox before driver provisioning.
crates/openshell-server/src/compute/mod.rs:522	Compute runtime invokes CreateSandbox on the selected driver after persistence.
crates/openshell-driver-kubernetes/src/config.rs:187	KubernetesComputeConfig has no warm-pool, claim, template, or extension API configuration today.
crates/openshell-driver-kubernetes/src/driver.rs:84	Kubernetes driver only defines agents.x-k8s.io Sandbox group/version constants.
crates/openshell-driver-kubernetes/src/driver.rs:273	agent_sandbox_api constructs dynamic APIs only for the core Sandbox kind.
crates/openshell-driver-kubernetes/src/driver.rs:280	Runtime API discovery is cached for core Sandbox versions; extension API discovery would need a parallel path.
crates/openshell-driver-kubernetes/src/driver.rs:528	create_sandbox always creates a direct Sandbox CR today.
crates/openshell-driver-kubernetes/src/driver.rs:647	delete_sandbox deletes a direct Sandbox; claim mode should likely delete the claim.
crates/openshell-driver-kubernetes/src/driver.rs:710	Watch path watches direct Sandbox resources and Kubernetes events; claim mode must correlate claims and adopted sandboxes.
crates/openshell-driver-kubernetes/src/driver.rs:857	OpenShell labels direct Sandbox CRs with the gateway sandbox ID and managed-by labels.
crates/openshell-driver-kubernetes/src/driver.rs:882	Driver converts observed DynamicObject Sandbox resources into internal driver sandbox observations.
crates/openshell-driver-kubernetes/src/driver.rs:1436	sandbox_to_k8s_spec renders final direct Sandbox spec and always includes default workspace volumeClaimTemplates.
crates/openshell-driver-kubernetes/src/driver.rs:1527	Pod labels can carry OpenShell sandbox ID only because the pod template is rendered at create time.
crates/openshell-driver-kubernetes/src/driver.rs:1549	Pod annotation openshell.io/sandbox-id is injected at pod creation and is used by Kubernetes token bootstrap validation.
crates/openshell-driver-kubernetes/src/driver.rs:1775	Workspace persistence injection is tied to the generated pod template and default volumeClaimTemplates.
crates/openshell-driver-kubernetes/src/driver.rs:2125	Status parsing expects core Sandbox status fields including agentPod.
deploy/helm/openshell/templates/role.yaml:12	Helm RBAC only grants agents.x-k8s.io sandboxes and sandboxes/status today.
tasks/scripts/helm-k3s-local.sh:142	Local k3s setup applies only upstream manifest.yaml, even though the script comment mentions extensions.
docs/reference/sandbox-compute-drivers.mdx:313	Published docs state that the Kubernetes driver creates agents.x-k8s.io Sandbox resources directly.
docs/kubernetes/setup.mdx:31	Kubernetes setup docs instruct users to install only the core Agent Sandbox manifest.

Current Behavior

OpenShell creates each Kubernetes sandbox as a fresh Agent Sandbox Sandbox CR. Every create pays for pod scheduling, image pull if needed, supervisor sideload setup, workspace PVC provisioning/seeding, service account token projection, and supervisor startup.

The driver does not create or consume extensions.agents.x-k8s.io SandboxTemplate, SandboxWarmPool, or SandboxClaim. The Helm chart does not grant permissions on those resources, and the local k3s helper does not install extensions.yaml.

OpenShell identity is currently early-bound. The Kubernetes driver renders the OpenShell sandbox ID/name into pod environment and pod annotations before the Agent Sandbox controller creates the pod. IssueSandboxToken then validates the pod and its controlling Sandbox owner reference against that identity. A real warm-pool adoption path breaks that assumption because the warm pod may already be running before OpenShell knows which request will claim it.

What Would Need to Change

The Kubernetes driver needs an extension-resource API layer for extensions.agents.x-k8s.io, likely mirroring the existing dynamic AgentSandboxApi pattern. It should detect the served extension API version, construct Api<DynamicObject> handles for SandboxClaim, SandboxTemplate, and possibly SandboxWarmPool, and preserve the existing Kubernetes API timeout behavior.

The create path needs a claim mode. When warm-pool mode is enabled, create_sandbox should create a SandboxClaim that references an existing warm pool rather than directly creating agents.x-k8s.io/Sandbox. It must still attach enough labels/annotations to the claim and/or adopted Sandbox for OpenShell to resolve the public sandbox ID, name, and lifecycle.

The read/delete/watch paths need claim-aware mapping. get_sandbox, list_sandboxes, delete_sandbox, sandbox_exists, and watch_sandboxes currently assume the OpenShell sandbox name is the core Sandbox CR name. Claim mode must decide whether the claim name is the OpenShell sandbox name, how to resolve the adopted Sandbox name from claim status, how to map pod events to OpenShell sandbox IDs, and whether deleting a public sandbox deletes the claim, the adopted Sandbox, or both.

The biggest design change is late-bound identity. OpenShell must avoid putting per-request env/PVC data on the SandboxClaim if it wants true warm adoption, because upstream marks those claims as cold-start paths. Candidate directions include a supervisor late-bind mode, an entrypoint wrapper that waits for claim/adoption metadata before launching openshell-sandbox, or an upstream extension point that starts warm pods suspended until claim metadata is available. This decision also affects IssueSandboxToken, relay startup, provider environment injection, policy sync, and initial command launch.

The operator/config surface needs a decision. Possible surfaces are a driver-level default warm pool in [openshell.drivers.kubernetes], a per-create driver_config.kubernetes.warm_pool_ref, a Helm-managed pool configuration, or a public CLI/API concept. A narrow first pass should prefer driver-owned config unless maintainers want warm pools to become a cross-driver product concept.

Helm and local development must add optional extension support. RBAC needs verbs for extension resources when the gateway creates claims/templates/pools. Local k3s and e2e setup need to apply upstream extensions.yaml for warm-pool tests. Docs must explain that core Agent Sandbox installation is still enough for cold mode, while warm pools require extensions.

Alternative Approaches Considered

1. Create SandboxClaim with per-claim env and PVC templates. This is the smallest code change, but upstream explicitly forces cold start when env or volumeClaimTemplates are set. It would add API compatibility without solving the launch-latency goal.

2. Operator-managed SandboxTemplate and SandboxWarmPool, OpenShell-created claims. This aligns best with upstream and keeps pool sizing/rollouts in Kubernetes, but requires late-bound OpenShell identity and careful lifecycle mapping.

3. OpenShell-managed pools using only core Sandbox resources. This avoids the extension API and gives OpenShell full control, but duplicates upstream Agent Sandbox warm-pool controllers and increases gateway responsibility. This is not recommended unless upstream semantics cannot support OpenShell identity binding.

4. Base-template-only warm pools as a first milestone. This could warm only stable images, supervisor sideload, and workspace setup while deferring dynamic env/providers/uploads/GPU/runtime-class support. It may be useful, but the issue should be explicit about which create features remain cold paths.

Patterns to Follow

Use driver-owned configuration for Kubernetes-only fields. The public API and compute-driver proto already keep platform-specific fields behind driver_config, and ComputeRuntime selects only the active driver's block before calling the driver.

Preserve gateway persistence before Kubernetes provisioning. The current create path stores the public sandbox record before driver create because token bootstrap, status reconciliation, and cleanup rely on that record existing.

Follow the existing dynamic CRD discovery pattern instead of introducing compile-time typed Kubernetes clients. The Kubernetes driver already probes Agent Sandbox API versions at runtime, caches the result, and uses DynamicObject.

Keep OpenShell-managed labels and annotations as the primary reconciliation surface. The driver already labels direct Sandbox resources with the sandbox ID and managed-by labels and annotates pods with openshell.io/sandbox-id.

Keep tests close to the rendering and lifecycle boundaries. The Kubernetes driver already has focused JSON rendering tests for pod templates, workspace persistence, supervisor sideload, driver config parsing, and status conversion.

Proposed Approach

Add Kubernetes warm-pool support as an optional claim-based provisioning mode in the Kubernetes driver, backed by upstream Agent Sandbox extension CRDs. Start by modeling an operator-provided warm pool reference, probably through driver_config.kubernetes.warm_pool_ref and/or [openshell.drivers.kubernetes], and keep direct Sandbox creation as the default. Before claiming performance wins, solve and document late-bound sandbox identity so pre-warmed pods can safely become a specific OpenShell sandbox without forcing upstream cold-start fallback. Treat SandboxTemplate and SandboxWarmPool lifecycle ownership as an explicit product decision: either operator-managed at first, or Helm-managed only when the chart owns pool settings.

Scope Assessment

• Complexity: High
• Confidence: Medium - upstream Agent Sandbox is the right primitive, but late-bound identity and product surface need human decisions.
• Estimated files to change: 12-20
• Issue type: feat

Risks & Open Questions

• Should OpenShell expose warm pools as Kubernetes-only driver config, Helm/operator defaults, or a user-facing sandbox create feature?
• How should OpenShell bind sandbox ID/name, relay endpoint, and bootstrap metadata to a pod that may already be running?
• Which create-time features should be compatible with warm adoption: per-sandbox env, providers, policy, uploads, GPU, runtime class, workspace PVC sizing, and custom image?
• If per-claim env or PVC templates force upstream cold start, should OpenShell reject those combinations in warm-pool mode or silently fall back to direct/cold provisioning?
• Should OpenShell create SandboxTemplate and SandboxWarmPool, or only create SandboxClaim against operator-managed pools?
• How should pool template rollouts handle OpenShell supervisor image updates, sandbox image updates, policy bootstrap changes, and workspace persistence changes?
• What is the canonical OpenShell logical resource in claim mode: the claim, the adopted Sandbox, or both?
• Failure cleanup must avoid leaking persisted OpenShell sandbox records, claims, adopted Sandboxes, pods, and pool replacement resources.
• RBAC expands to extension resources if the gateway creates claims/templates/pools.
• There is no direct LSM impact for CRD routing alone, but any late-bound supervisor startup or UID/GID/AppArmor changes must be rechecked on AppArmor, SELinux/OpenShift, and user-namespace clusters.

Test Considerations

• Add Kubernetes driver unit tests for extension API resource construction and version detection.
• Add JSON rendering tests for SandboxClaim, SandboxTemplate, and optional SandboxWarmPool objects if OpenShell owns any of them.
• Add tests for direct-create vs claim-create branching based on driver config and gateway config.
• Add tests for get_sandbox, list_sandboxes, delete_sandbox, sandbox_exists, and watch_sandboxes mapping in claim mode.
• Add server compute tests to verify driver config pass-through and persistence cleanup behavior still works when driver create fails after claim creation.
• Add Helm tests for new values, rendered gateway.toml, and RBAC on extensions.agents.x-k8s.io.
• Add a Kubernetes e2e path that installs upstream extensions.yaml, creates a small warm pool, creates multiple OpenShell sandboxes, and verifies at least one claim adopts a pre-warmed Sandbox.
• Update tasks/scripts/helm-k3s-local.sh or add a dedicated e2e task so extension CRDs are installed only when warm-pool tests require them.
• Update docs: docs/reference/gateway-config.mdx, docs/reference/sandbox-compute-drivers.mdx, docs/kubernetes/setup.mdx, Helm README/values, and crates/openshell-driver-kubernetes/README.md.
• Update architecture/sandbox.md if late-bound supervisor identity changes sandbox startup invariants.

---

Created by spike investigation. Use build-from-issue to plan and implement.

[drew]

Related: #1447