fix(sdk): scale the token refresh lead to the token lifetime

crates/basilica-sdk/src/auth/simple_manager.rs

@@ -7,7 +7,6 @@ use super::refresh::refresh_access_token;
 use super::token_store::TokenStore;
 use super::types::{get_sdk_data_dir, AuthError, AuthMethod, AuthResult, TokenSet};
 use std::sync::Arc;
-use std::time::Duration;
 use tokio::sync::Mutex;
 use tracing::{debug, info};
 
@@ -19,9 +18,6 @@ pub struct TokenManager {
 }
 
 impl TokenManager {
-    /// Pre-emptive refresh threshold (60 minutes before expiry)
-    const REFRESH_THRESHOLD: Duration = Duration::from_secs(3600);
-
     /// Create a new token manager with direct tokens
     pub fn new_direct(access_token: String, refresh_token: String) -> Self {
         let tokens = TokenSet::new(access_token, refresh_token);
@@ -108,12 +104,10 @@ impl TokenManager {
     }
 
     /// Check if token should be refreshed
+    ///
+    /// Defers to `TokenSet`, so this manager and the CLI renew on the same
+    /// schedule instead of each carrying its own threshold.
     fn should_refresh(&self, token_set: &TokenSet) -> bool {
-        if token_set.is_expired() {
-            return true;
-        }
-
-        // Pre-emptive refresh if expiring within threshold
-        token_set.expires_within(Self::REFRESH_THRESHOLD)
+        token_set.is_expired() || token_set.needs_refresh()
     }
 }

crates/basilica-sdk/src/auth/types.rs

@@ -7,7 +7,13 @@ use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
 use etcetera::{choose_base_strategy, BaseStrategy};
 use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
-use std::time::{SystemTime, UNIX_EPOCH};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+/// Ceiling on how early a token is renewed, however long its lifetime.
+const MAX_REFRESH_LEAD: Duration = Duration::from_secs(60 * 60);
+
+/// Lead used when the token carries no `iat`, so its lifetime is unknown.
+const FALLBACK_REFRESH_LEAD: Duration = Duration::from_secs(5 * 60);
 
 /// Result type for authentication operations
 pub type AuthResult<T> = Result<T, AuthError>;
@@ -51,9 +57,8 @@ impl TokenSet {
         }
     }
 
-    /// Extract expiration from JWT token
-    /// Returns the exp claim from the JWT if it can be decoded
-    fn decode_jwt_exp(token: &str) -> Option<u64> {
+    /// Read a numeric claim from the access token's payload.
+    fn decode_jwt_claim(token: &str, claim: &str) -> Option<u64> {
         // JWT has three parts: header.payload.signature
         let parts: Vec<&str> = token.split('.').collect();
         if parts.len() != 3 {
@@ -66,15 +71,41 @@ impl TokenSet {
         // Decode base64url without padding (JWT uses base64url encoding)
         let decoded = URL_SAFE_NO_PAD.decode(payload).ok()?;
 
-        // Parse JSON and extract exp claim
+        // Parse JSON and extract the requested claim
         let json: serde_json::Value = serde_json::from_slice(&decoded).ok()?;
-        json.get("exp")?.as_u64()
+        json.get(claim)?.as_u64()
     }
 
     /// Get the expiration time by decoding JWT
     fn get_expiration(&self) -> Option<u64> {
         // Always decode from JWT token
-        Self::decode_jwt_exp(&self.access_token)
+        Self::decode_jwt_claim(&self.access_token, "exp")
+    }
+
+    /// Total lifetime the issuer gave this token, from its own claims.
+    fn lifetime(&self) -> Option<Duration> {
+        let exp = Self::decode_jwt_claim(&self.access_token, "exp")?;
+        let iat = Self::decode_jwt_claim(&self.access_token, "iat")?;
+        exp.checked_sub(iat).map(Duration::from_secs)
+    }
+
+    /// How long before expiry a refresh should happen.
+    ///
+    /// Half the token's own lifetime, capped at an hour. A fixed lead is wrong
+    /// for short-lived tokens: with a one-hour token, a fixed one-hour lead is
+    /// satisfied the instant the token is minted, so every single call would
+    /// refresh. Scaling with the lifetime keeps roughly half the token usable
+    /// before the first renewal, whatever the issuer configures.
+    ///
+    /// A 24-hour token still gets the one-hour lead it had before this was
+    /// derived, so long-lived tokens behave exactly as they always did.
+    fn refresh_lead(&self) -> Duration {
+        match self.lifetime() {
+            Some(lifetime) => MAX_REFRESH_LEAD.min(lifetime / 2),
+            // No `iat` to measure against; fall back to the same small buffer
+            // the token stores use rather than assuming a long lifetime.
+            None => FALLBACK_REFRESH_LEAD,
+        }
     }
 
     /// Check if the access token is expired
@@ -91,13 +122,15 @@ impl TokenSet {
         }
     }
 
-    /// Check if the token needs refresh (expires within 60 minutes)
+    /// Check whether the token is close enough to expiry to renew.
+    ///
+    /// The window scales with the token's own lifetime; see [`Self::refresh_lead`].
     pub fn needs_refresh(&self) -> bool {
-        self.expires_within(std::time::Duration::from_secs(60 * 60))
+        self.expires_within(self.refresh_lead())
     }
 
     /// Check if the token expires within the specified duration
-    pub fn expires_within(&self, duration: std::time::Duration) -> bool {
+    pub fn expires_within(&self, duration: Duration) -> bool {
         match self.get_expiration() {
             Some(expires_at) => {
                 let now = SystemTime::now()
@@ -112,17 +145,17 @@ impl TokenSet {
     }
 
     /// Get time until token expiration
-    pub fn time_until_expiry(&self) -> Option<std::time::Duration> {
+    pub fn time_until_expiry(&self) -> Option<Duration> {
         match self.get_expiration() {
             Some(expires_at) => {
                 let now = SystemTime::now()
                     .duration_since(UNIX_EPOCH)
                     .unwrap()
                     .as_secs();
                 if expires_at > now {
-                    Some(std::time::Duration::from_secs(expires_at - now))
+                    Some(Duration::from_secs(expires_at - now))
                 } else {
-                    Some(std::time::Duration::from_secs(0)) // Already expired
+                    Some(Duration::from_secs(0)) // Already expired
                 }
             }
             None => None, // No expiration time
@@ -219,3 +252,72 @@ pub fn get_sdk_data_dir() -> AuthResult<PathBuf> {
     // Use the same path as the CLI for consistency
     Ok(strategy.data_dir().join("basilica"))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Build an access token whose payload carries the given claims. Only the
+    /// payload is read, so the header and signature can be anything.
+    fn token_with(iat: u64, exp: u64) -> TokenSet {
+        let payload = URL_SAFE_NO_PAD.encode(format!(r#"{{"iat":{iat},"exp":{exp}}}"#));
+        TokenSet::new(format!("header.{payload}.signature"), "refresh".to_string())
+    }
+
+    fn now() -> u64 {
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_secs()
+    }
+
+    #[test]
+    fn lead_is_half_the_lifetime_for_a_short_token() {
+        // One-hour token: renew in the final half hour, not immediately.
+        let t = token_with(now(), now() + 3600);
+        assert_eq!(t.refresh_lead(), Duration::from_secs(1800));
+    }
+
+    #[test]
+    fn a_freshly_minted_short_token_does_not_need_refresh() {
+        // The bug this replaces: a fixed one-hour lead made every call to a
+        // one-hour token refresh, because it was always "expiring within" it.
+        let t = token_with(now(), now() + 3600);
+        assert!(!t.needs_refresh());
+    }
+
+    #[test]
+    fn a_short_token_past_halfway_needs_refresh() {
+        let t = token_with(now() - 1900, now() + 1700);
+        assert!(t.needs_refresh());
+    }
+
+    #[test]
+    fn long_lived_tokens_keep_the_one_hour_lead() {
+        // A 24-hour token behaves exactly as it did before the lead was derived.
+        let t = token_with(now(), now() + 86_400);
+        assert_eq!(t.refresh_lead(), Duration::from_secs(3600));
+        assert!(!t.needs_refresh());
+    }
+
+    #[test]
+    fn a_long_token_inside_the_last_hour_needs_refresh() {
+        let t = token_with(now() - 84_000, now() + 2_400);
+        assert!(t.needs_refresh());
+    }
+
+    #[test]
+    fn without_iat_the_lead_falls_back_to_the_small_buffer() {
+        let payload = URL_SAFE_NO_PAD.encode(format!(r#"{{"exp":{}}}"#, now() + 3600));
+        let t = TokenSet::new(format!("header.{payload}.signature"), "refresh".to_string());
+        assert_eq!(t.refresh_lead(), FALLBACK_REFRESH_LEAD);
+        assert!(!t.needs_refresh());
+    }
+
+    #[test]
+    fn an_expired_token_needs_refresh() {
+        let t = token_with(now() - 7200, now() - 3600);
+        assert!(t.is_expired());
+        assert!(t.needs_refresh());
+    }
+}