feat(server): keep pool sandbox sessions on their supervisor's gateway

.config/nextest.toml

@@ -32,11 +32,11 @@ kubernetes-shared-supervisor-pool = { max-threads = 1 }
 filter = "test(/gateway_(scale_and_rollout|pod_rolls)$/)"
 test-group = "kubernetes-ha"
 
-# Both change the shared supervisor pool of the release they run against: one
-# kills a member, the other scales the pool. Never run them at once. Test
-# names, not binary(), for the reason above.
+# These change the shared supervisor pool of the release they run against:
+# they kill members or gateways, or scale the pool. Never run them at once.
+# Test names, not binary(), for the reason above.
 [[profile.e2e-kubernetes.overrides]]
-filter = "test(/^(losing_a_supervisor_pod_moves_its_sandboxes_without_restarting_them|scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running)$/)"
+filter = "test(/^(losing_a_supervisor_pod_moves_its_sandboxes_without_restarting_them|scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running|pool_sessions_follow_their_supervisors_gateway)$/)"
 test-group = "kubernetes-shared-supervisor-pool"
 
 # Scaling the pool runs two bounded Helm upgrades (up to 6 minutes each), waits
@@ -46,6 +46,11 @@ test-group = "kubernetes-shared-supervisor-pool"
 filter = "test(/^scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running$/)"
 slow-timeout = { period = "60s", terminate-after = 20 }
 
+# Three bounded rollouts (up to 5 minutes each) and four settle and exec rounds.
+[[profile.e2e-kubernetes.overrides]]
+filter = "test(/^pool_sessions_follow_their_supervisors_gateway$/)"
+slow-timeout = { period = "60s", terminate-after = 20 }
+
 # Relative to the profile store dir (`e2e/rust/target/nextest/e2e-kubernetes/`).
 [profile.e2e-kubernetes.junit]
 path = "../../../../../results/e2e-kubernetes.xml"

TESTING.md

@@ -672,7 +672,7 @@ Available task variants:
 | `e2e:kubernetes:v1alpha1` | Agent Sandbox v1alpha1 compatibility |
 | `e2e:kubernetes:external-driver` | External Kubernetes driver sidecar |
 | `e2e:kubernetes:shared-supervisor` | One-member shared supervisor pool for the gateway's sandboxes (gateway TLS through the port-forward) |
-| `e2e:kubernetes:shared-supervisor-ha` | Two-replica shared supervisor pool: kills the member that owns a sandbox and checks the sandbox is adopted, not restarted |
+| `e2e:kubernetes:shared-supervisor-ha` | Two-replica shared supervisor pool behind two gateway replicas on PostgreSQL: kills the member that owns a sandbox and checks the sandbox is adopted, not restarted; crashes and restarts gateways and checks each sandbox's session stays on its supervisor's gateway |
 | `e2e:kubernetes:shared-supervisor-scaling` | One-member shared supervisor pool: scales it to two members and back with `helm upgrade`, checking sandboxes stay running and the removed member hands its sandboxes off |
 
 Kubernetes e2e environment variables:

crates/openshell-server/src/grpc/supervisor_host.rs

@@ -59,7 +59,16 @@ impl SupervisorHost for SupervisorHostService {
                 Status::failed_precondition("this gateway does not host a shared supervisor")
             })?
             .clone();
+        // A stopping replica sheds its hosts so they reattach to a live one
+        // before their sandbox sessions follow them there.
+        let mut shutdown = self.state.supervisor_sessions.subscribe_shutdown();
+        if *shutdown.borrow() {
+            return Err(Status::unavailable("gateway is shutting down"));
+        }
         let outbound = controller.attach(identity, request.into_inner()).await?;
+        let outbound = futures::StreamExt::take_until(outbound, async move {
+            let _ = shutdown.wait_for(|shutdown| *shutdown).await;
+        });
         Ok(Response::new(Box::pin(outbound)))
     }
 }

crates/openshell-server/src/lib.rs

@@ -836,6 +836,11 @@ pub(crate) async fn run_server(
             "shared supervisor assignment stream enabled; minting supervisor tokens"
         );
         info!("shared supervisor pool enabled; reconciling pool membership and ownership");
+        // Sandbox sessions dial this endpoint with their tokens.
+        if let Some(peer_endpoint) = state.peer_endpoint.as_deref() {
+            validate_peer_endpoint_scheme(&state.config, peer_endpoint)?;
+        }
+        controller.set_peer_endpoint(state.peer_endpoint.clone().unwrap_or_default());
         supervisor_pool::SupervisorPool::new(
             state.store.clone(),
             controller.clone(),

crates/openshell-server/src/supervisor_session.rs

@@ -435,6 +435,11 @@ impl SupervisorSessionRegistry {
         self.admission_closed.store(true, Ordering::Release);
     }
 
+    /// Watch for [`Self::shutdown`] beginning.
+    pub(crate) fn subscribe_shutdown(&self) -> watch::Receiver<bool> {
+        self.shutdown.subscribe()
+    }
+
     /// Close control sessions and wait for tracked ownership cleanup. Call
     /// after compute shutdown so supervisors can finish normal stop reporting.
     pub(crate) async fn shutdown(&self, timeout: Duration) -> Result<(), String> {

crates/openshell-supervisor-control/README.md

@@ -46,7 +46,9 @@ re-implementing the control contracts.
   stream attaches, ends, or announces leaving, so the gateway refreshes
   membership records at once. The gateway closes a member's stream after
   `pool_timing::GATEWAY_SILENCE_TIMEOUT` of silence, and drops a member that
-  stops reading as soon.
+  stops reading as soon. `HostAccepted` carries this replica's peer endpoint
+  (`set_peer_endpoint`); the host dials its instances' sandbox sessions
+  there, so each session runs on the replica serving its supervisor.
 - `pool_timing`: the timing contract between pool hosts and the gateway: the
   15 s membership TTL, the 2 s record refresh, 1 s heartbeats, host silence
   detection and bounded reconnect attempts, and `reconnect_budget`, the worst

crates/openshell-supervisor-control/src/controller.rs

@@ -356,6 +356,8 @@ struct Inner {
     /// stream attached or ended, or a member announced leaving.
     membership: watch::Sender<u64>,
     observed: Mutex<Observed>,
+    /// This replica's peer endpoint, sent in every `HostAccepted`.
+    peer_endpoint: Mutex<String>,
 }
 
 /// Gateway-owned desired assignments and observed status of the shared
@@ -406,10 +408,17 @@ impl SupervisorHostController {
                 revision: watch::Sender::new(0),
                 membership: watch::Sender::new(0),
                 observed: Mutex::new(Observed::default()),
+                peer_endpoint: Mutex::new(String::new()),
             }),
         }
     }
 
+    /// Tell attaching hosts to dial their sandbox sessions to `peer_endpoint`.
+    /// Call it before any host attaches.
+    pub fn set_peer_endpoint(&self, peer_endpoint: String) {
+        *lock(&self.inner.peer_endpoint) = peer_endpoint;
+    }
+
     /// Offer every instance's runtime `recovery_timeout_secs` instead of
     /// [`POOL_RECOVERY_TIMEOUT_SECS`]. Call it before any host attaches.
     pub fn set_recovery_timeout_secs(&self, recovery_timeout_secs: u32) {
@@ -843,13 +852,15 @@ impl SupervisorHostController {
         let (tx, rx) = mpsc::channel(OUTBOUND_QUEUE);
         let expires_at = identity.expires_at;
         let (session_id, superseded) = self.begin_session(identity, &hello)?;
+        let peer_endpoint = lock(&self.inner.peer_endpoint).clone();
         tx.send(Ok(HostGatewayMessage {
             payload: Some(host_gateway_message::Payload::Accepted(HostAccepted {
                 session_id: session_id.to_string(),
                 heartbeat_interval: Some(prost_types::Duration {
                     seconds: i64::try_from(heartbeat.as_secs()).unwrap_or(15),
                     nanos: 0,
                 }),
+                peer_endpoint,
             })),
         }))
         .await

crates/openshell-supervisor-process/src/delegated.rs

@@ -93,6 +93,7 @@ pub async fn start_boundary_access(
     config_apply_tx: tokio::sync::mpsc::Sender<crate::supervisor_session::ConfigApplyRequest>,
     config_push_updates: tokio::sync::watch::Sender<bool>,
     host_key: Option<russh::keys::PrivateKey>,
+    gateway_home: Option<tokio::sync::watch::Receiver<Option<String>>>,
 ) -> Result<BoundaryAccess> {
     let instance_id = uuid::Uuid::new_v4().to_string();
     let terminating = Arc::new(AtomicBool::new(false));
@@ -130,6 +131,7 @@ pub async fn start_boundary_access(
                     session_id_updates: supervisor_session_updates,
                     config_apply_tx,
                     config_push_updates,
+                    gateway_home,
                 },
             );
             (Some(task), Some(accepted))
@@ -374,6 +376,7 @@ mod tests {
             tokio::sync::mpsc::channel(1).0,
             tokio::sync::watch::channel(false).0,
             None,
+            None,
         )
         .await
         .expect("access without SSH is portable");
@@ -415,6 +418,7 @@ mod tests {
             tokio::sync::mpsc::channel(1).0,
             tokio::sync::watch::channel(false).0,
             None,
+            None,
         )
         .await;
         assert!(

crates/openshell-supervisor-process/src/supervisor_session.rs

@@ -44,6 +44,9 @@ use openshell_core::transport_errors::is_expected_transport_close_status;
 
 const INITIAL_BACKOFF: Duration = Duration::from_secs(1);
 const MAX_BACKOFF: Duration = Duration::from_secs(30);
+/// How long a shared host's session waits for the host to attach to a
+/// gateway replica before dialing another address.
+const GATEWAY_HOME_WAIT: Duration = Duration::from_secs(5);
 const CONFIG_APPLY_TIMEOUT: Duration = Duration::from_mins(1);
 
 /// A stream-delivered configuration update awaiting application by the
@@ -129,6 +132,9 @@ pub struct SessionRuntimeContext {
     /// Publishes whether the current session pushes configuration, so the
     /// runtime stops polling while it does.
     pub config_push_updates: watch::Sender<bool>,
+    /// Peer endpoint of the gateway replica serving this supervisor's shared
+    /// host. When set, sessions dial it instead of the configured address.
+    pub gateway_home: Option<watch::Receiver<Option<String>>>,
 }
 
 /// Parse a gRPC endpoint URI into an OCSF `Endpoint` (host + port). Falls back
@@ -423,6 +429,7 @@ pub fn spawn_with_readiness(
         session_id_updates: runtime.session_id_updates,
         config_apply_tx: runtime.config_apply_tx,
         config_push_updates: runtime.config_push_updates,
+        gateway_home: runtime.gateway_home,
         ready_tx,
     };
     (tokio::spawn(run_session_loop(config)), ready_rx)
@@ -442,6 +449,7 @@ struct SessionConfig {
     /// Publishes whether the current session pushes configuration. A
     /// reconnect to a gateway that does not push resumes polling.
     config_push_updates: watch::Sender<bool>,
+    gateway_home: Option<watch::Receiver<Option<String>>>,
     ready_tx: watch::Sender<bool>,
 }
 
@@ -455,9 +463,22 @@ async fn run_session_loop(config: SessionConfig) {
     // back to the configured address, which load-balances across replicas.
     let mut target = config.gateway.clone();
     let mut redirected = false;
+    // A shared host's sessions dial the replica serving the host, as
+    // redirected so that replica serves them, unless they are following a
+    // redirect away from it.
+    let mut gateway_home = config.gateway_home.clone();
+    let mut follow_home = true;
 
     loop {
         attempt += 1;
+        if follow_home
+            && let Some(home) = gateway_home.as_mut()
+            && let Some(endpoint) = wait_for_home(home, None).await
+        {
+            target = config.gateway.with_url(endpoint);
+            redirected = true;
+        }
+        follow_home = true;
 
         let result = run_single_session(&config, &target, redirected, attempt).await;
         if let Some(updates) = &config.session_id_updates {
@@ -485,10 +506,27 @@ async fn run_session_loop(config: SessionConfig) {
                     peer_endpoint = %peer_endpoint,
                     "supervisor session: following gateway redirect"
                 );
-                target = config.gateway.with_url(peer_endpoint);
+                // Prefer the replica the shared host moved to.
+                let home = match gateway_home.as_mut() {
+                    Some(home) => wait_for_home(home, Some(target.url())).await,
+                    None => None,
+                };
+                target = config.gateway.with_url(home.unwrap_or(peer_endpoint));
                 redirected = true;
+                follow_home = false;
                 // No backoff: this is an expected handoff, not a failure.
             }
+            Ok(SessionOutcome::HomeMoved { peer_endpoint }) => {
+                backoff.on_redirect(config.ready_tx.send_replace(false));
+                info!(
+                    sandbox_id = %config.sandbox_id,
+                    peer_endpoint = %peer_endpoint,
+                    "supervisor session: following the shared host to another gateway"
+                );
+                target = config.gateway.with_url(peer_endpoint);
+                redirected = true;
+                follow_home = false;
+            }
             Err(e) => {
                 let accepted = config.ready_tx.send_replace(false);
                 let event = session_failed_event(
@@ -555,6 +593,44 @@ impl ReconnectBackoff {
     }
 }
 
+/// The peer endpoint of the replica serving the shared host, once it is one
+/// other than `avoid`, waiting up to [`GATEWAY_HOME_WAIT`] while no replica
+/// serves the host. `None` at once when the serving replica has no endpoint.
+async fn wait_for_home(
+    home: &mut watch::Receiver<Option<String>>,
+    avoid: Option<&str>,
+) -> Option<String> {
+    tokio::time::timeout(
+        GATEWAY_HOME_WAIT,
+        home.wait_for(|home| home.as_deref().is_some_and(|home| Some(home) != avoid)),
+    )
+    .await
+    .ok()?
+    .ok()
+    .and_then(|home| (*home).clone())
+    .filter(|home| !home.is_empty())
+}
+
+/// The peer endpoint of the shared host, once it is one other than `current`,
+/// including a move that landed before this is first polled. Pending forever
+/// without a shared host.
+async fn home_moved(home: Option<&mut watch::Receiver<Option<String>>>, current: &str) -> String {
+    let Some(home) = home else {
+        return std::future::pending().await;
+    };
+    loop {
+        let endpoint = home.borrow_and_update().clone();
+        if let Some(endpoint) =
+            endpoint.filter(|endpoint| !endpoint.is_empty() && endpoint != current)
+        {
+            return endpoint;
+        }
+        if home.changed().await.is_err() {
+            return std::future::pending().await;
+        }
+    }
+}
+
 /// Keep asking not to be redirected when a redirected attempt never got a
 /// session, so a stale ring pointing at a dead replica cannot bounce us back
 /// to it. A session that was accepted and later dropped starts placement over.
@@ -579,6 +655,9 @@ enum SessionOutcome {
         peer_endpoint: String,
         owner_replica_id: String,
     },
+    /// The shared host reattached to another replica. The caller reconnects
+    /// there.
+    HomeMoved { peer_endpoint: String },
 }
 
 async fn run_single_session(
@@ -693,9 +772,13 @@ async fn run_single_session(
     // Main loop: receive gateway messages + send heartbeats.
     let mut heartbeat_interval = tokio::time::interval(Duration::from_secs(heartbeat_secs));
     heartbeat_interval.tick().await; // skip immediate tick
+    let mut home = config.gateway_home.clone();
 
     loop {
         tokio::select! {
+            peer_endpoint = home_moved(home.as_mut(), target.url()) => {
+                return Ok(SessionOutcome::HomeMoved { peer_endpoint });
+            }
             msg = inbound.message() => {
                 let msg = match map_session_stream_message(
                     msg,
@@ -1636,6 +1719,77 @@ mod ocsf_event_tests {
         assert!(!redirect_survives_failure(false, true));
     }
 
+    #[tokio::test(start_paused = true)]
+    async fn shared_host_sessions_wait_for_the_hosts_gateway() {
+        let (tx, mut home) = watch::channel(None);
+        assert_eq!(wait_for_home(&mut home, None).await, None);
+
+        tx.send_replace(Some("https://gw-a:8443".to_string()));
+        assert_eq!(
+            wait_for_home(&mut home, None).await.as_deref(),
+            Some("https://gw-a:8443")
+        );
+        assert_eq!(
+            wait_for_home(&mut home, Some("https://gw-a:8443")).await,
+            None
+        );
+
+        let mut moving = home.clone();
+        let moved =
+            tokio::spawn(
+                async move { wait_for_home(&mut moving, Some("https://gw-a:8443")).await },
+            );
+        tokio::task::yield_now().await;
+        tx.send_replace(Some("https://gw-b:8443".to_string()));
+        assert_eq!(moved.await.unwrap().as_deref(), Some("https://gw-b:8443"));
+    }
+
+    #[tokio::test]
+    async fn a_host_gateway_without_a_peer_endpoint_is_not_waited_for() {
+        let (_tx, mut home) = watch::channel(Some(String::new()));
+        let started = tokio::time::Instant::now();
+        assert_eq!(wait_for_home(&mut home, None).await, None);
+        assert!(started.elapsed() < GATEWAY_HOME_WAIT);
+
+        let moved = tokio::time::timeout(
+            Duration::from_millis(50),
+            home_moved(Some(&mut home), "https://gw-a:8443"),
+        )
+        .await;
+        assert!(moved.is_err(), "an endpointless gateway is not a move");
+    }
+
+    #[tokio::test]
+    async fn a_host_move_during_session_setup_is_followed() {
+        let (tx, mut home) = watch::channel(Some("https://gw-a:8443".to_string()));
+        // The session dialed A; the host moved to B before it was accepted,
+        // and the move was already observed.
+        tx.send_replace(Some("https://gw-b:8443".to_string()));
+        home.mark_unchanged();
+        assert_eq!(
+            home_moved(Some(&mut home), "https://gw-a:8443").await,
+            "https://gw-b:8443"
+        );
+    }
+
+    #[tokio::test]
+    async fn an_established_session_follows_its_host_to_another_gateway() {
+        let (tx, mut home) = watch::channel(Some("https://gw-a:8443".to_string()));
+        let moved =
+            tokio::spawn(async move { home_moved(Some(&mut home), "https://gw-a:8443").await });
+
+        tx.send_replace(None);
+        tx.send_replace(Some("https://gw-a:8443".to_string()));
+        tokio::task::yield_now().await;
+        assert!(
+            !moved.is_finished(),
+            "a reattach to the same gateway is not a move"
+        );
+
+        tx.send_replace(Some("https://gw-b:8443".to_string()));
+        assert_eq!(moved.await.unwrap(), "https://gw-b:8443");
+    }
+
     #[test]
     fn reconnect_backoff_doubles_until_a_session_is_accepted() {
         let mut backoff = ReconnectBackoff::new();