fix(sdk): refresh OIDC tokens on raw routes and harden rotation

Raw gRPC access never triggered OIDC refresh: a client that only used
raw_grpc/raw_inference kept sending the initial bearer until it expired,
with no proactive or reactive refresh. Add raw_grpc_fresh and
raw_inference_fresh accessors that refresh before returning the client,
plus force_refresh for reactive recovery after an Unauthenticated raw
RPC.

Guard the single-flight refresh commit against a concurrent replace().
The in-flight attempt now records the generation it started from and
skips its write when an external replace() has advanced it, so timer or
callback driven rotation is no longer clobbered by a slower refresh.

Remove TokenSource::snapshot(): it returned an empty string under write
contention and had no consumer on the CLI/TUI path. Tests read committed
state directly instead.

Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>

Author: maxdubrinsky

crates/openshell-sdk/README.md

@@ -15,6 +15,13 @@ Designed in [RFC 0008](../../rfc/0008-shared-sdk-core-and-ts-binding/README.md).
   surface doesn't yet cover (inference, providers, policy, logs, settings, SSH,
   forwarding).
 
+The curated surface drives OIDC refresh automatically (proactively before a
+request and reactively on `Unauthenticated`). The plain `raw_grpc`/
+`raw_inference` accessors do not: they return a client bound to the current
+token. When a refresher is wired, use `raw_grpc_fresh`/`raw_inference_fresh`
+for a proactive refresh before the call, and `force_refresh` to recover after
+a raw RPC returns `Unauthenticated`.
+
 ## Responsibilities
 
 - Construct the gRPC channel and select the transport (plaintext vs TLS).

crates/openshell-sdk/src/client.rs

@@ -83,21 +83,61 @@ impl OpenShellClient {
     ///
     /// Use this when the curated surface below doesn't expose the RPC or
     /// field you need.
+    ///
+    /// This does **not** drive OIDC refresh: it returns a client bound to
+    /// the interceptor's current bearer slot without checking expiry or
+    /// retrying on `Unauthenticated`. A client that only ever issues raw
+    /// RPCs keeps sending the initial token until it expires. When a
+    /// refresher is wired, prefer [`OpenShellClient::raw_grpc_fresh`] (or
+    /// interleave a curated call) so rotation reaches the shared slot, and
+    /// call [`OpenShellClient::force_refresh`] to recover on a rejected
+    /// token.
     pub fn raw_grpc(&self) -> AuthedGrpcClient {
         proto::open_shell_client::OpenShellClient::with_interceptor(
             self.channel.clone(),
             self.interceptor.clone(),
         )
     }
 
+    /// Like [`OpenShellClient::raw_grpc`], but proactively refreshes the
+    /// bearer token first when a refresher is wired and the token is within
+    /// the refresh skew of expiry. The returned client reads the same live
+    /// slot, so the refreshed token applies to every RPC issued through it.
+    ///
+    /// Reactive retry on `Unauthenticated` remains the caller's
+    /// responsibility for raw RPCs: on a rejected token, call
+    /// [`OpenShellClient::force_refresh`] and reissue.
+    pub async fn raw_grpc_fresh(&self) -> Result<AuthedGrpcClient> {
+        self.ensure_fresh().await?;
+        Ok(self.raw_grpc())
+    }
+
     /// Authenticated gRPC client for the inference service.
+    ///
+    /// Like [`OpenShellClient::raw_grpc`], this does not drive OIDC refresh;
+    /// use [`OpenShellClient::raw_inference_fresh`] when a refresher is wired.
     pub fn raw_inference(&self) -> AuthedInferenceClient {
         proto::inference_client::InferenceClient::with_interceptor(
             self.channel.clone(),
             self.interceptor.clone(),
         )
     }
 
+    /// Like [`OpenShellClient::raw_inference`], but proactively refreshes the
+    /// bearer token first (see [`OpenShellClient::raw_grpc_fresh`]).
+    pub async fn raw_inference_fresh(&self) -> Result<AuthedInferenceClient> {
+        self.ensure_fresh().await?;
+        Ok(self.raw_inference())
+    }
+
+    /// Force an OIDC refresh and write the new token into the live bearer
+    /// slot, regardless of expiry. Returns `true` when a refresher is wired
+    /// and a fresh token was minted, `false` for static auth. Use after a
+    /// raw RPC returns `Unauthenticated` to recover before reissuing it.
+    pub async fn force_refresh(&self) -> Result<bool> {
+        self.refresh_on_unauthorized().await
+    }
+
     /// Gateway health snapshot.
     pub async fn health(&self) -> Result<Health> {
         let resp = self
@@ -451,7 +491,110 @@ fn map_status(status: tonic::Status) -> SdkError {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::refresh::{Refresh, RefreshError};
+    use std::sync::atomic::{AtomicUsize, Ordering};
     use std::sync::{Arc, RwLock};
+    use std::time::{SystemTime, UNIX_EPOCH};
+    use tonic::transport::Channel;
+
+    struct StubRefresher {
+        calls: Arc<AtomicUsize>,
+    }
+
+    #[async_trait::async_trait]
+    impl Refresh for StubRefresher {
+        async fn refresh(&self) -> std::result::Result<RefreshedToken, RefreshError> {
+            let n = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
+            Ok(RefreshedToken::new(format!("token-{n}")).with_expires_at(
+                SystemTime::now()
+                    .duration_since(UNIX_EPOCH)
+                    .unwrap()
+                    .as_secs()
+                    + 3600,
+            ))
+        }
+    }
+
+    /// Build an OIDC client wired to a refresher, with a near-expiry initial
+    /// token, over a lazy channel that never actually connects (no RPC is
+    /// issued in these tests).
+    fn oidc_client_with_refresher(calls: Arc<AtomicUsize>) -> OpenShellClient {
+        let interceptor = EdgeAuthInterceptor::new(Some("initial"), None).unwrap();
+        let bearer_slot = interceptor.bearer_slot();
+        let near = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_secs()
+            + 5;
+        let source = TokenSource::new(
+            RefreshedToken::new("initial").with_expires_at(near),
+            Arc::new(StubRefresher { calls }),
+        );
+        let channel = Channel::from_static("http://127.0.0.1:1").connect_lazy();
+        OpenShellClient {
+            channel,
+            interceptor,
+            token_source: Some(source),
+            bearer_slot,
+        }
+    }
+
+    fn slot_token(slot: &BearerSlot) -> String {
+        slot.read()
+            .unwrap()
+            .clone()
+            .unwrap()
+            .to_str()
+            .unwrap()
+            .to_string()
+    }
+
+    #[tokio::test]
+    async fn raw_grpc_does_not_refresh() {
+        // Regression (P1): the plain raw accessor must not be relied on for
+        // rotation — it hands back a client bound to the current token.
+        let calls = Arc::new(AtomicUsize::new(0));
+        let client = oidc_client_with_refresher(Arc::clone(&calls));
+        let _raw = client.raw_grpc();
+        assert_eq!(calls.load(Ordering::SeqCst), 0, "raw_grpc must not refresh");
+        assert_eq!(
+            slot_token(client.bearer_slot.as_ref().unwrap()),
+            "Bearer initial"
+        );
+    }
+
+    #[tokio::test]
+    async fn raw_grpc_fresh_refreshes_near_expiry() {
+        // Regression (P1): the _fresh accessor proactively rotates a
+        // near-expiry token into the shared slot before returning a client.
+        let calls = Arc::new(AtomicUsize::new(0));
+        let client = oidc_client_with_refresher(Arc::clone(&calls));
+        let _raw = client.raw_grpc_fresh().await.unwrap();
+        assert_eq!(
+            calls.load(Ordering::SeqCst),
+            1,
+            "raw_grpc_fresh must refresh a near-expiry token"
+        );
+        assert_eq!(
+            slot_token(client.bearer_slot.as_ref().unwrap()),
+            "Bearer token-1",
+            "refreshed token must reach the live slot"
+        );
+    }
+
+    #[tokio::test]
+    async fn force_refresh_rotates_and_reports_wired() {
+        // Regression (P1): reactive recovery path for raw callers.
+        let calls = Arc::new(AtomicUsize::new(0));
+        let client = oidc_client_with_refresher(Arc::clone(&calls));
+        let refreshed = client.force_refresh().await.unwrap();
+        assert!(refreshed, "force_refresh reports a wired refresher");
+        assert_eq!(calls.load(Ordering::SeqCst), 1);
+        assert_eq!(
+            slot_token(client.bearer_slot.as_ref().unwrap()),
+            "Bearer token-1"
+        );
+    }
 
     #[test]
     fn store_bearer_rejects_malformed_token_and_keeps_previous() {

crates/openshell-sdk/src/refresh.rs

@@ -177,15 +177,6 @@ impl TokenSource {
         }
     }
 
-    /// Current token without checking expiry. Used by the sync gRPC
-    /// interceptor, which can't await.
-    pub fn snapshot(&self) -> String {
-        self.state
-            .try_read()
-            .map(|s| s.token.clone())
-            .unwrap_or_default()
-    }
-
     /// Async-fetch the current token, refreshing if it's within `skew` of
     /// expiry. Single-flight: concurrent callers share one refresh.
     ///
@@ -246,14 +237,25 @@ impl TokenSource {
                 let state = Arc::clone(&self.state);
                 let flight_slot = Arc::clone(&self.flight);
                 let epoch = flight.epoch.wrapping_add(1);
+                // Generation this attempt refreshes *from*. If it advances
+                // while the refresh is in flight, an external `replace()`
+                // installed a newer token and this result must not clobber it.
+                let start_generation = expected_generation;
                 let future: RefreshFuture = async move {
                     let outcome = match refresher.refresh().await {
                         Ok(token) => {
                             let mut state = state.write().await;
-                            state.token.clone_from(&token.access_token);
-                            state.expires_at = token.expires_at;
-                            state.generation = state.generation.wrapping_add(1);
-                            Ok(token.access_token)
+                            if state.generation == start_generation {
+                                state.token.clone_from(&token.access_token);
+                                state.expires_at = token.expires_at;
+                                state.generation = state.generation.wrapping_add(1);
+                                Ok(token.access_token)
+                            } else {
+                                // A concurrent `replace()` (or newer token)
+                                // landed mid-flight; honor it rather than
+                                // regressing to this now-stale result.
+                                Ok(state.token.clone())
+                            }
                         }
                         Err(err) => Err(SdkError::from(err).to_string()),
                     };
@@ -373,6 +375,12 @@ mod tests {
     /// Drive one leader into the refresher, then queue `followers` more
     /// callers that must join the in-flight attempt rather than start their
     /// own. Returns the refresher call count and per-caller outcomes.
+    /// Read the committed token straight from state, bypassing any refresh
+    /// logic. Test-only inspection helper.
+    async fn state_token(source: &TokenSource) -> String {
+        source.state.read().await.token.clone()
+    }
+
     async fn run_coalesced(fail: bool) -> (usize, Vec<Result<String>>) {
         let calls = Arc::new(AtomicUsize::new(0));
         let entered = Arc::new(tokio::sync::Notify::new());
@@ -551,7 +559,64 @@ mod tests {
         let token = source.refresh_now().await.unwrap();
         assert_eq!(token, "token-1", "forced refresh must mint a new token");
         assert_eq!(calls.load(Ordering::SeqCst), 1);
-        assert_eq!(source.snapshot(), "token-1", "slot must observe new token");
+        assert_eq!(
+            state_token(&source).await,
+            "token-1",
+            "state must observe new token"
+        );
+    }
+
+    #[tokio::test]
+    async fn in_flight_refresh_preserves_external_replace() {
+        // Regression (P2b): a `replace()` that lands while a refresh is in
+        // flight must survive. The in-flight refresh started from an older
+        // generation, so its result is discarded in favor of the externally
+        // installed token instead of clobbering it.
+        let calls = Arc::new(AtomicUsize::new(0));
+        let entered = Arc::new(tokio::sync::Notify::new());
+        let release = Arc::new(tokio::sync::Notify::new());
+        let refresher = Arc::new(GatedRefresher {
+            calls: Arc::clone(&calls),
+            entered: Arc::clone(&entered),
+            release: Arc::clone(&release),
+            fail: false,
+        });
+        let source = TokenSource::new(RefreshedToken::new("initial").with_expires_at(0), refresher);
+
+        // Leader starts a refresh and blocks inside refresher.refresh().
+        let leader = {
+            let src = source.clone();
+            tokio::spawn(async move { src.refresh_now().await })
+        };
+        entered.notified().await;
+
+        // External rotation installs an authoritative token mid-flight.
+        source
+            .replace(
+                RefreshedToken::new("external-rotation").with_expires_at(
+                    SystemTime::now()
+                        .duration_since(UNIX_EPOCH)
+                        .unwrap()
+                        .as_secs()
+                        + 3600,
+                ),
+            )
+            .await;
+
+        // Release the in-flight refresh; it must yield to the replacement.
+        release.notify_waiters();
+        let leader_token = leader.await.unwrap().unwrap();
+
+        assert_eq!(
+            leader_token, "external-rotation",
+            "leader must observe the external replacement, not its own result"
+        );
+        assert_eq!(
+            state_token(&source).await,
+            "external-rotation",
+            "replace() must survive the in-flight refresh"
+        );
+        assert_eq!(calls.load(Ordering::SeqCst), 1);
     }
 
     #[tokio::test]

crates/openshell-sdk/tests/client_mock.rs

@@ -878,3 +878,58 @@ async fn unauthenticated_without_refresher_surfaces_error() {
         "exactly one attempt; no retry without a refresher"
     );
 }
+
+/// Regression (P1): raw RPCs do not drive refresh. A call through the plain
+/// `raw_grpc()` accessor keeps sending the seeded token, while
+/// `raw_grpc_fresh()` proactively refreshes a near-expiry token into the
+/// shared slot so the subsequent raw call is accepted.
+#[tokio::test]
+async fn raw_grpc_fresh_refreshes_before_raw_call() {
+    let state = Arc::new(MockState {
+        require_bearer: Some("Bearer fresh-token".to_string()),
+        ..Default::default()
+    });
+    let endpoint = start_mock(state.clone()).await;
+
+    let calls = Arc::new(AtomicU32::new(0));
+    let refresher = Arc::new(OneShotRefresher {
+        calls: Arc::clone(&calls),
+    });
+
+    // Near-expiry token so the proactive path will refresh it.
+    let near = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap()
+        .as_secs()
+        + 5;
+    let mut config = ClientConfig::new(&endpoint);
+    config.auth = Some(AuthConfig::Oidc {
+        token: "stale-token".to_string(),
+        expires_at: Some(near),
+        refresh: Some(refresher),
+    });
+    let client = OpenShellClient::connect(config).await.unwrap();
+
+    // Plain raw accessor does not refresh: the seeded token is rejected.
+    let err = client
+        .raw_grpc()
+        .health(proto::HealthRequest {})
+        .await
+        .unwrap_err();
+    assert_eq!(err.code(), tonic::Code::Unauthenticated);
+    assert_eq!(
+        calls.load(Ordering::SeqCst),
+        0,
+        "raw_grpc must not trigger a refresh"
+    );
+
+    // _fresh accessor refreshes the near-expiry token first; the raw call is
+    // then accepted with the fresh bearer.
+    let mut grpc = client.raw_grpc_fresh().await.unwrap();
+    grpc.health(proto::HealthRequest {}).await.unwrap();
+    assert_eq!(
+        calls.load(Ordering::SeqCst),
+        1,
+        "raw_grpc_fresh must refresh the near-expiry token exactly once"
+    );
+}