api: ask the substrate for capabilities instead of matching on it

crates/libs/api/src/controllers/sandbox/lifecycle.rs

@@ -1,4 +1,4 @@
 use pplx_space_common_core::types::SandboxId;
 
 use super::*;
 use crate::controllers::db_contention_error;
@@ -361,11 +361,13 @@
 }
 
 impl PublishedSuspendFence {
-    fn for_substrate(substrate: &crate::controllers::substrate::Substrate) -> Self {
-        match substrate {
-            crate::controllers::substrate::Substrate::Direct(_) => Self::DirectRestoreBaseline,
-            #[cfg(feature = "kubernetes")]
-            crate::controllers::substrate::Substrate::Kubernetes(_) => Self::Enforced,
+    fn for_capabilities(
+        capabilities: crate::controllers::substrate::SubstrateCapabilities,
+    ) -> Self {
+        if capabilities.caller_completes_claimed_restore {
+            Self::Enforced
+        } else {
+            Self::DirectRestoreBaseline
         }
     }
 }
@@ -375,10 +377,13 @@
 /// A suspend claim can land after the node finishes waking but before this
 /// control-plane write. In that case the node-side wake succeeded, so callers
 /// may continue, but this stale completion must not replace `suspending` with
-/// `started`. Kubernetes also fences a finalized `suspended` state. Direct
-/// restores retain `suspended` as their owned baseline because that substrate
-/// does not claim `restoring` before its node call. Other CAS losses retain the
-/// existing error behavior.
+/// `started`. A substrate whose restore-shaped wakes claim `restoring` before
+/// the node call and leave `restoring -> started` to this CAS (Kubernetes)
+/// also fences a finalized `suspended` state. Direct does not: its
+/// worker-local wake leaves its pre-wake lifecycle (`suspended`, or
+/// `restoring` kept by an earlier lease restore) for this CAS to advance, and
+/// its lease restore writes `started` before returning.
+/// Other CAS losses retain the existing error behavior.
 async fn complete_wake_lifecycle(
     pool: &PgPool,
     sandbox_id: SandboxId,
@@ -1659,7 +1664,9 @@
                     if sandbox
                         .lifecycle
                         .parse::<SandboxLifecycle>()
-                        .is_ok_and(|lc| wake_completion::shield_applies(&self.substrate, lc))
+                        .is_ok_and(|lc| {
+                            wake_completion::shield_applies(self.substrate.capabilities(), lc)
+                        })
                     {
                         completion_guard.arm(wake_completion::wake_completion_redrive(
                             pool.clone(),
@@ -1735,7 +1742,7 @@
                     // completion must never overwrite if a stop, delete, or
                     // newer suspend wins while substrate work is in flight.
                     let published_suspend_fence =
-                        PublishedSuspendFence::for_substrate(&self.substrate);
+                        PublishedSuspendFence::for_capabilities(self.substrate.capabilities());
                     let cas_result =
                         complete_wake_lifecycle(pool, id, published_suspend_fence).await;
                     // Explicit exit either way from here (Ok, cas-lost, or DB
@@ -1826,7 +1833,9 @@
                     if sandbox
                         .lifecycle
                         .parse::<SandboxLifecycle>()
-                        .is_ok_and(|lc| wake_completion::shield_applies(&self.substrate, lc))
+                        .is_ok_and(|lc| {
+                            wake_completion::shield_applies(self.substrate.capabilities(), lc)
+                        })
                     {
                         completion_guard.arm(wake_completion::wake_completion_redrive(
                             pool.clone(),
@@ -1875,20 +1884,21 @@
                     // Refuse to overwrite a terminal, teardown, or newer
                     // suspend lifecycle.
                     // We can't guard on `lifecycle = 'restoring'` alone
-                    // because the two substrates enter restore from different
+                    // because the two substrates reach this CAS from different
                     // lifecycles: the Kubernetes substrate's
                     // `begin_restore_lifecycle` flips
                     // `suspended -> restoring` before the substrate call, so
                     // the row sits in `restoring` during the slow part; the
-                    // direct substrate's `restore_execution` doesn't touch
-                    // lifecycle at all, so the row sits in `suspended` for the
-                    // duration. The terminal-state exclusion is the only
+                    // direct substrate's `restore_execution` claims
+                    // `restoring` under its own lease and writes `started`
+                    // itself before returning, so the row is already
+                    // `started` here. The terminal-state exclusion is the only
                     // predicate that catches "sweeper landed" uniformly across
                     // substrates. On race-loss we surface a 5xx so the caller
                     // re-issues rather than silently observing a `started` row
                     // pointing at a deleted execution.
                     let published_suspend_fence =
-                        PublishedSuspendFence::for_substrate(&self.substrate);
+                        PublishedSuspendFence::for_capabilities(self.substrate.capabilities());
                     let cas_result =
                         complete_wake_lifecycle(pool, id, published_suspend_fence).await;
                     // Explicit exit either way from here — see the Wake arm.

crates/libs/api/src/controllers/sandbox/lifecycle/wake_completion_postgres_tests.rs

@@ -126,3 +126,47 @@ async fn wake_completion_cas_outcome_for_every_prior_lifecycle() -> TestResult {
     }
     Ok(())
 }
+
+/// The wake and restore arms pick the fence from the running substrate. On a
+/// row still `suspended` at completion, Kubernetes, whose restore leaves
+/// `restoring -> started` to this CAS, defers to it as a newer suspend; Direct,
+/// whose worker-local wake keeps `suspended` as its baseline, advances it.
+#[tokio::test(flavor = "current_thread")]
+async fn wake_completion_fence_follows_the_running_substrate() -> TestResult {
+    use crate::controllers::substrate::{Substrate, direct::DirectSubstrate};
+
+    let db = pg_test_db().await?;
+    let pool = &db.pool;
+
+    #[cfg_attr(not(feature = "kubernetes"), allow(unused_mut))]
+    let mut cases = vec![(
+        Substrate::Direct(Box::new(DirectSubstrate::new(pool.clone()))),
+        "started",
+    )];
+    #[cfg(feature = "kubernetes")]
+    cases.push((
+        Substrate::Kubernetes(Box::new(
+            crate::controllers::substrate::kubernetes::tests::substrate_with_store(pool.clone()),
+        )),
+        "suspended",
+    ));
+
+    for (substrate, expected) in cases {
+        let sandbox_id = SandboxId::new();
+        insert_sandbox(pool, sandbox_id, "suspended").await?;
+
+        let fence = PublishedSuspendFence::for_capabilities(substrate.capabilities());
+        complete_wake_lifecycle(pool, sandbox_id, fence).await?;
+
+        let lifecycle = lifecycle_and_placement(pool, sandbox_id)
+            .await?
+            .map(|(lifecycle, _)| lifecycle);
+        assert_eq!(
+            lifecycle.as_deref(),
+            Some(expected),
+            "{} substrate",
+            substrate.managed_by()
+        );
+    }
+    Ok(())
+}

crates/libs/api/src/controllers/sandbox/mod.rs

@@ -1122,8 +1122,8 @@ impl SandboxController {
         self.substrate.managed_by()
     }
 
-    pub fn supports_node_pool_placement(&self) -> bool {
-        self.substrate.supports_node_pool_placement()
+    pub fn substrate_capabilities(&self) -> super::substrate::SubstrateCapabilities {
+        self.substrate.capabilities()
     }
 
     pub async fn set_node_schedulability(

crates/libs/api/src/controllers/sandbox/queries.rs

@@ -1,10 +1,10 @@
 use pplx_space_common_core::types::{OrgId, SandboxId, UserId};
 
 use super::*;
 use crate::controllers::substrate::{
     AdmissionPoolClaim, AdmittedPoolClaim, AdmittedPoolClaimLookup, CreateBootMode,
     PooledClaimRefusal, PooledExecutionClaim, PooledExecutionClaimRequest,
-    PooledExecutionPoolShape, PooledPrepareFailureCleanup,
+    PooledExecutionPoolShape, PooledPrepareFailureCleanup, SubstrateCapabilities,
 };
 use crate::controllers::{db_contention_error, internal_error, pool_store_error};
 use crate::reconcile_queue::finalize_sandbox_db_delete_tx;
@@ -566,7 +566,7 @@
 }
 
 struct ColdCreateShapeInput<'a> {
-    managed_by: &'static str,
+    capabilities: SubstrateCapabilities,
     runtime: RuntimeKind,
     boot_mode: CreateBootMode,
     template: &'a TemplateName,
@@ -594,7 +594,7 @@
     input: ColdCreateShapeInput<'_>,
 ) -> Result<(SandboxResources, DiskIoLimits), ControllerError> {
     let ColdCreateShapeInput {
-        managed_by,
+        capabilities,
         runtime,
         boot_mode,
         template,
@@ -603,7 +603,7 @@
         requested,
         restoring,
     } = input;
-    if managed_by != "k8s" || runtime != RuntimeKind::Firecracker {
+    if !capabilities.resolves_cold_create_shape || runtime != RuntimeKind::Firecracker {
         return Ok((requested, DiskIoLimits::default()));
     }
     if boot_mode == CreateBootMode::Fallthrough
@@ -1318,7 +1318,7 @@
         let (resources, disk_io_limits) = resolve_cold_create_resources(
             pool,
             ColdCreateShapeInput {
-                managed_by: self.substrate.managed_by(),
+                capabilities: self.substrate.capabilities(),
                 runtime,
                 boot_mode,
                 template: &template_name,

crates/libs/api/src/controllers/sandbox/queries/create_postgres_tests.rs

@@ -716,6 +716,10 @@ impl SubstrateAdapter for RetryingPoolSubstrate {
         "k8s"
     }
 
+    fn capabilities(&self) -> SubstrateCapabilities {
+        SubstrateCapabilities::KUBERNETES
+    }
+
     async fn create_execution(
         &self,
         _request: CreateExecutionRequest<'_>,
@@ -1048,6 +1052,10 @@ impl SubstrateAdapter for NodeRpcStubbedKubernetes {
         self.inner.managed_by()
     }
 
+    fn capabilities(&self) -> SubstrateCapabilities {
+        self.inner.capabilities()
+    }
+
     async fn create_execution(
         &self,
         _request: CreateExecutionRequest<'_>,
@@ -1756,11 +1764,13 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap
     let release_id = insert_template_release(&db.pool, "app-runtime", "sha256:shape").await?;
     insert_inventory_shape(&db.pool, "builder", "sha256:shape", (4, 4096), 0).await?;
     let template = TemplateName::parse("app-runtime").unwrap();
-    let resolve_with = |managed_by, boot_mode, requested, restoring| {
+    let k8s = SubstrateCapabilities::KUBERNETES;
+    let space = SubstrateCapabilities::DIRECT;
+    let resolve_with = |substrate, boot_mode, requested, restoring| {
         resolve_cold_create_resources(
             &db.pool,
             ColdCreateShapeInput {
-                managed_by,
+                capabilities: substrate,
                 runtime: RuntimeKind::Firecracker,
                 boot_mode,
                 template: &template,
@@ -1772,7 +1782,7 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap
         )
     };
     let resolve =
-        |managed_by, boot_mode, requested| resolve_with(managed_by, boot_mode, requested, false);
+        |substrate, boot_mode, requested| resolve_with(substrate, boot_mode, requested, false);
     let unset = || SandboxResources {
         vcpu: None,
         memory_mib: None,
@@ -1781,15 +1791,15 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap
     let shape =
         |resources: &SandboxResources| (resources.vcpu, resources.memory_mib, resources.disk_mib);
 
-    let (resources, _) = resolve("k8s", CreateBootMode::Fallthrough, unset()).await?;
+    let (resources, _) = resolve(k8s, CreateBootMode::Fallthrough, unset()).await?;
     assert_eq!(
         shape(&resources),
         (Some(4), Some(4096), None),
         "an unshaped Kubernetes create must not boot at the shim's fixed defaults"
     );
 
     let (resources, _) = resolve(
-        "k8s",
+        k8s,
         CreateBootMode::Cold,
         SandboxResources {
             vcpu: None,
@@ -1804,14 +1814,14 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap
         "requested values win; only unset ones come from the template"
     );
 
-    let (resources, _) = resolve("space", CreateBootMode::Fallthrough, unset()).await?;
+    let (resources, _) = resolve(space, CreateBootMode::Fallthrough, unset()).await?;
     assert_eq!(
         shape(&resources),
         (None, None, None),
         "direct creates keep the node's snapshot or metadata shape"
     );
 
-    let (resources, _) = resolve_with("k8s", CreateBootMode::Cold, unset(), true).await?;
+    let (resources, _) = resolve_with(k8s, CreateBootMode::Cold, unset(), true).await?;
     assert_eq!(
         shape(&resources),
         (None, None, None),
@@ -1829,7 +1839,7 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap
     .bind(release_id)
     .execute(&db.pool)
     .await?;
-    let (resources, _) = resolve("k8s", CreateBootMode::Fallthrough, unset()).await?;
+    let (resources, _) = resolve(k8s, CreateBootMode::Fallthrough, unset()).await?;
     assert_eq!(
         shape(&resources),
         (Some(2), Some(8192), Some(32768)),

crates/libs/api/src/controllers/sandbox/wake_completion.rs

@@ -41,7 +41,7 @@ use pplx_space_sandbox_store as sandbox_store;
 use sqlx::PgPool;
 use tokio::sync::Semaphore;
 
-use crate::controllers::substrate::Substrate;
+use crate::controllers::substrate::{Substrate, SubstrateCapabilities};
 
 /// How often the redrive re-checks the row and re-attempts hydration. The
 /// node-side restore this waits on takes seconds to low tens of seconds;
@@ -78,25 +78,27 @@ pub(super) fn redrive_budget() -> Duration {
 }
 
 /// Whether the cancellation shield applies to this wake: only restore-shaped
-/// wakes (`suspended`/`restoring` pre-wake lifecycle) on the Kubernetes
-/// substrate. Everything else keeps existing semantics:
-/// * the direct substrate's restore never flips the row out of `suspended`
-///   mid-flight, so a cancelled direct restore leaves an already-recoverable
-///   state;
+/// wakes (`suspended`/`restoring` pre-wake lifecycle) on a substrate that
+/// leaves `restoring -> started` to the caller's completion CAS (Kubernetes).
+/// Everything else keeps existing semantics:
+/// * the direct substrate does not guarantee this: its worker-local
+///   wake leaves the pre-wake lifecycle (`suspended`, or `restoring` kept by
+///   an earlier lease restore) for the completion CAS, and its lease restore
+///   writes `started` itself, so a cancelled direct wake leaves `suspended`,
+///   `started`, or a `restoring` row that a later restore reclaims once any
+///   lease expires;
 /// * pause/resume and stop/start wakes don't take the restore claim and have
 ///   no execution to hydrate — their (much rarer) cancellation class is a
 ///   separate piece of work.
-pub(super) fn shield_applies(substrate: &Substrate, pre_wake_lifecycle: SandboxLifecycle) -> bool {
+pub(super) fn shield_applies(
+    capabilities: SubstrateCapabilities,
+    pre_wake_lifecycle: SandboxLifecycle,
+) -> bool {
     let restore_shaped = matches!(
         pre_wake_lifecycle,
         SandboxLifecycle::Suspended | SandboxLifecycle::Restoring
     );
-    let kubernetes = match substrate {
-        Substrate::Direct(_) => false,
-        #[cfg(feature = "kubernetes")]
-        Substrate::Kubernetes(_) => true,
-    };
-    restore_shaped && kubernetes
+    restore_shaped && capabilities.caller_completes_claimed_restore
 }
 
 /// Drop-guard that owns the boxed redrive future while the wake handler is

crates/libs/api/src/controllers/substrate/direct.rs

@@ -1,4 +1,4 @@
 use pplx_space_common_core::sql;
 use pplx_space_common_svc::sql::contention;
 use std::fmt;
 use std::time::Duration;
@@ -27,7 +27,8 @@
     ExecutionBinding, NodeApiTarget, NodeInventoryEntry, OwnedNodeApiTarget, PooledExecutionClaim,
     PooledExecutionClaimAttempt, PooledExecutionClaimRequest, PooledExecutionPoolShape,
     RetireExecutionOutcome, RetireExecutionRequest, RuntimeClient, SubstrateAdapter,
-    SubstrateError, connect_builder, connect_worker, node_api_endpoint, substrate_db_error,
+    SubstrateCapabilities, SubstrateError, connect_builder, connect_worker, node_api_endpoint,
+    substrate_db_error,
 };
 use crate::controllers::auto_lifecycle::policy_to_proto;
 use crate::controllers::{ControllerError, db_contention_error};
@@ -3087,6 +3088,10 @@
         "space"
     }
 
+    fn capabilities(&self) -> SubstrateCapabilities {
+        SubstrateCapabilities::DIRECT
+    }
+
     fn pool_eligible_for_request(
         &self,
         boot_mode: CreateBootMode,

crates/libs/api/src/controllers/substrate/kubernetes/adapter/substrate_adapter.rs

@@ -1,4 +1,4 @@
 //! `SubstrateAdapter` implementation for `KubernetesSubstrate`.
 
 use super::*;
 
@@ -7,8 +7,8 @@
         "k8s"
     }
 
-    fn supports_node_pool_placement(&self) -> bool {
-        true
+    fn capabilities(&self) -> SubstrateCapabilities {
+        SubstrateCapabilities::KUBERNETES
     }
 
     async fn create_execution<'a>(

crates/libs/api/src/controllers/substrate/kubernetes/mod.rs

@@ -83,7 +83,8 @@ use super::{
     ExecutionIdentity, NODE_RPC_TIMEOUT, NodeApiTarget, NodeInventoryEntry, ObservedPod,
     ObservedPodIdentity, OwnedNodeApiTarget, PodFenceMismatch, PodIdentityFence, RecordedExecution,
     RetireExecutionOutcome, RetireExecutionRequest, RetireFence, RuntimeClient, SubstrateAdapter,
-    SubstrateError, TeardownProgress, WaitReadyTimeoutPhase, connect_worker, node_api_endpoint,
+    SubstrateCapabilities, SubstrateError, TeardownProgress, WaitReadyTimeoutPhase, connect_worker,
+    node_api_endpoint,
 };
 
 const START_WAIT_FAST_POLL_INTERVAL: Duration = Duration::from_millis(250);