feat(sdk): route Rust SDK exec to the sandbox owner replica

crates/openshell-sdk/README.md

@@ -62,6 +62,11 @@ endpoint. Authorization is stripped by default. Select `BearerPassthrough` only
 when the sandbox application validates its own bearer credential. The returned
 `SandboxRef::service_urls` map contains each routed URL under the same name.
 
+`exec` sends the sandbox's owner replica in the `x-openshell-replica` header so
+an edge proxy with replica routing can skip the peer relay. The owner is cached
+for 30 seconds, and a call is retried once without the header if the owner's
+route has no ready endpoints.
+
 Curated calls without a workspace argument explicitly select the `default`
 workspace. Cross-workspace listing uses the separate `*_all_workspaces`
 methods and requires Platform Admin access.

crates/openshell-sdk/src/client.rs

@@ -14,6 +14,7 @@ use crate::error::{Result, SdkError};
 use crate::pagination::{Page, Pager};
 use crate::raw::AuthedGrpcClient;
 use crate::refresh::{RefreshedToken, TokenSource};
+use crate::routing::{OwnerCache, owner_replica, retry_unrouted, routed_to, with_owner_hint};
 use crate::types::{
     DeleteOptions, DeletionResult, ExecOptions, ExecResult, Health, ListOptions, SandboxPhase,
     SandboxRef, SandboxSpec, SandboxTemplateCreateSpec, SandboxTemplateListOptions,
@@ -26,6 +27,7 @@ use std::collections::HashMap;
 use std::future::Future;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
+use tonic::metadata::AsciiMetadataValue;
 use tonic::transport::Channel;
 
 /// Async client for a single `OpenShell` gateway.
@@ -43,6 +45,7 @@ pub struct OpenShellClient {
     /// Live bearer slot the interceptor reads; refreshed tokens are written
     /// here so rotation reaches in-flight requests. `None` for non-OIDC auth.
     bearer_slot: Option<BearerSlot>,
+    owners: OwnerCache,
 }
 
 impl OpenShellClient {
@@ -60,6 +63,7 @@ impl OpenShellClient {
             interceptor,
             token_source,
             bearer_slot,
+            owners: OwnerCache::default(),
         })
     }
 
@@ -75,6 +79,7 @@ impl OpenShellClient {
             interceptor,
             token_source: None,
             bearer_slot,
+            owners: OwnerCache::default(),
         }
     }
 
@@ -625,16 +630,7 @@ impl OpenShellClient {
             no_login_shell: opts.no_login_shell,
         };
 
-        // Open the stream under the same OIDC-aware auth policy as unary RPCs
-        // (proactive refresh, then one reactive retry on `Unauthenticated`).
-        // Mid-stream rotation is out of scope; streaming retry is tracked
-        // separately.
-        let mut stream = self
-            .unary(|mut grpc| {
-                let request = request.clone();
-                async move { grpc.exec_sandbox(request).await }
-            })
-            .await?;
+        let mut stream = self.open_exec("default", request).await?;
 
         let mut stdout = Vec::new();
         let mut stderr = Vec::new();
@@ -786,19 +782,27 @@ impl OpenShellClient {
     /// response, force a refresh and retry exactly once. No-op auth behaves
     /// as a plain single call.
     async fn unary<T, F, Fut>(&self, call: F) -> Result<T>
+    where
+        F: Fn(AuthedGrpcClient) -> Fut,
+        Fut: Future<Output = std::result::Result<tonic::Response<T>, tonic::Status>>,
+    {
+        self.unary_response(call)
+            .await
+            .map(tonic::Response::into_inner)
+    }
+
+    /// [`OpenShellClient::unary`] that keeps the response metadata.
+    async fn unary_response<T, F, Fut>(&self, call: F) -> Result<tonic::Response<T>>
     where
         F: Fn(AuthedGrpcClient) -> Fut,
         Fut: Future<Output = std::result::Result<tonic::Response<T>, tonic::Status>>,
     {
         self.ensure_fresh().await?;
         match call(self.raw_grpc()).await {
-            Ok(resp) => Ok(resp.into_inner()),
+            Ok(resp) => Ok(resp),
             Err(status) if status.code() == tonic::Code::Unauthenticated => {
                 if self.refresh_on_unauthorized().await? {
-                    call(self.raw_grpc())
-                        .await
-                        .map(tonic::Response::into_inner)
-                        .map_err(map_status)
+                    call(self.raw_grpc()).await.map_err(map_status)
                 } else {
                     Err(map_status(status))
                 }
@@ -807,6 +811,53 @@ impl OpenShellClient {
         }
     }
 
+    /// Best effort: a failed lookup leaves the call unrouted.
+    async fn owner_of(&self, workspace: &str, name: &str) -> Option<AsciiMetadataValue> {
+        if let Some(cached) = self.owners.get(workspace, name) {
+            return cached.0;
+        }
+        let response = self
+            .unary_response(|mut grpc| {
+                let request = with_owner_hint(proto::GetSandboxRequest {
+                    name: name.to_string(),
+                    workspace_scope: Some(proto::workspace_selector(workspace)),
+                });
+                async move { grpc.get_sandbox(request).await }
+            })
+            .await
+            .ok()?;
+        let owner = owner_replica(&response);
+        self.owners.insert(workspace, name, owner.clone());
+        owner
+    }
+
+    /// Open `ExecSandbox` on the owner replica under the unary auth policy.
+    async fn open_exec(
+        &self,
+        workspace: &str,
+        request: proto::ExecSandboxRequest,
+    ) -> Result<tonic::Streaming<proto::ExecSandboxEvent>> {
+        let owner = self.owner_of(workspace, &request.sandbox).await;
+        let owners = &self.owners;
+        self.unary(|mut grpc| {
+            let request = request.clone();
+            let owner = owner.clone();
+            async move {
+                match grpc
+                    .exec_sandbox(routed_to(request.clone(), owner.as_ref()))
+                    .await
+                {
+                    Err(status) if retry_unrouted(owner.as_ref(), &status) => {
+                        owners.forget(workspace, &request.sandbox);
+                        grpc.exec_sandbox(request).await
+                    }
+                    response => response,
+                }
+            }
+        })
+        .await
+    }
+
     /// Proactive refresh: if a token source is wired and the token is within
     /// the refresh skew of expiry, mint a new one and store it in the live
     /// bearer slot. Tokens with no advertised expiry are left untouched.
@@ -1204,13 +1255,7 @@ impl WorkspaceScopedClient {
             no_login_shell: opts.no_login_shell,
         };
 
-        let mut stream = self
-            .client
-            .unary(|mut grpc| {
-                let request = request.clone();
-                async move { grpc.exec_sandbox(request).await }
-            })
-            .await?;
+        let mut stream = self.client.open_exec(&self.workspace, request).await?;
 
         let mut stdout = Vec::new();
         let mut stderr = Vec::new();
@@ -1547,6 +1592,7 @@ mod tests {
             interceptor,
             token_source: Some(source),
             bearer_slot,
+            owners: OwnerCache::default(),
         }
     }
 

crates/openshell-sdk/src/lib.rs

@@ -40,6 +40,7 @@ pub mod pagination;
 pub mod provider_readiness;
 pub mod raw;
 pub mod refresh;
+mod routing;
 pub mod transport;
 pub mod types;
 

crates/openshell-sdk/src/routing.rs

@@ -0,0 +1,121 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Routes long-lived sandbox calls to the replica that owns the sandbox.
+
+use openshell_core::replica_routing::{
+    OWNER_REPLICA_HEADER, OWNER_REQUEST_HEADER, ROUTE_REPLICA_HEADER,
+};
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex, PoisonError};
+use std::time::{Duration, Instant};
+use tonic::metadata::AsciiMetadataValue;
+
+/// A stale owner is safe: a live non-owner replica relays the call.
+const OWNER_TTL: Duration = Duration::from_secs(30);
+const MAX_OWNERS: usize = 1024;
+
+type Entries = HashMap<(String, String), (Instant, Option<AsciiMetadataValue>)>;
+
+/// Owner the gateway last reported; `None` means it reported none.
+pub struct Cached(pub Option<AsciiMetadataValue>);
+
+/// Reported owners keyed by workspace and sandbox name.
+#[derive(Clone, Default)]
+pub struct OwnerCache(Arc<Mutex<Entries>>);
+
+impl OwnerCache {
+    pub fn get(&self, workspace: &str, name: &str) -> Option<Cached> {
+        let entries = self.0.lock().unwrap_or_else(PoisonError::into_inner);
+        let (stored, owner) = entries.get(&key(workspace, name))?;
+        (stored.elapsed() < OWNER_TTL).then(|| Cached(owner.clone()))
+    }
+
+    pub fn insert(&self, workspace: &str, name: &str, owner: Option<AsciiMetadataValue>) {
+        let mut entries = self.0.lock().unwrap_or_else(PoisonError::into_inner);
+        if entries.len() >= MAX_OWNERS {
+            entries.clear();
+        }
+        entries.insert(key(workspace, name), (Instant::now(), owner));
+    }
+
+    pub fn forget(&self, workspace: &str, name: &str) {
+        self.0
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .remove(&key(workspace, name));
+    }
+}
+
+fn key(workspace: &str, name: &str) -> (String, String) {
+    (workspace.to_string(), name.to_string())
+}
+
+pub fn with_owner_hint<T>(message: T) -> tonic::Request<T> {
+    let mut request = tonic::Request::new(message);
+    request
+        .metadata_mut()
+        .insert(OWNER_REQUEST_HEADER, AsciiMetadataValue::from_static("1"));
+    request
+}
+
+pub fn owner_replica<T>(response: &tonic::Response<T>) -> Option<AsciiMetadataValue> {
+    response.metadata().get(OWNER_REPLICA_HEADER).cloned()
+}
+
+pub fn routed_to<T>(message: T, replica: Option<&AsciiMetadataValue>) -> tonic::Request<T> {
+    let mut request = tonic::Request::new(message);
+    if let Some(replica) = replica {
+        request
+            .metadata_mut()
+            .insert(ROUTE_REPLICA_HEADER, replica.clone());
+    }
+    request
+}
+
+/// Envoy Gateway answers a replica route with no ready endpoints with
+/// `grpc-status: 14` and no message; gateway errors always carry one.
+pub fn retry_unrouted(replica: Option<&AsciiMetadataValue>, status: &tonic::Status) -> bool {
+    let message = status.message();
+    replica.is_some()
+        && status.code() == tonic::Code::Unavailable
+        && (message.is_empty()
+            || message.contains("no healthy upstream")
+            || message.contains("connection failure"))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn only_unrouted_unavailable_is_retried() {
+        let replica = AsciiMetadataValue::from_static("openshell-1");
+        let unavailable = |message| tonic::Status::unavailable(message);
+        assert!(retry_unrouted(Some(&replica), &unavailable("")));
+        assert!(retry_unrouted(
+            Some(&replica),
+            &unavailable("no healthy upstream")
+        ));
+        assert!(!retry_unrouted(None, &unavailable("")));
+        assert!(!retry_unrouted(
+            Some(&replica),
+            &unavailable("supervisor not connected")
+        ));
+        assert!(!retry_unrouted(
+            Some(&replica),
+            &tonic::Status::not_found("")
+        ));
+    }
+
+    #[test]
+    fn cache_distinguishes_unknown_from_no_owner() {
+        let cache = OwnerCache::default();
+        assert!(cache.get("default", "box").is_none());
+        cache.insert("default", "box", None);
+        assert!(cache.get("default", "box").is_some_and(|c| c.0.is_none()));
+        assert!(cache.get("team", "box").is_none());
+        cache.forget("default", "box");
+        assert!(cache.get("default", "box").is_none());
+    }
+}