feat: implement telemetry authentication token management in FleetDeviceSession and related components

jwric · jwric · commit 00af202d83cc · 2026-03-02T17:14:25.000-05:00
diff --git a/crates/burn-central-fleet/src/error.rs b/crates/burn-central-fleet/src/error.rs
@@ -2,6 +2,8 @@ use crate::{model, state, telemetry};
 
 #[derive(Debug, thiserror::Error)]
 pub enum FleetError {
+    #[error("fleet registration failed: {0}")]
+    RegistrationFailed(String),
     #[error("fleet sync failed: {0}")]
     SyncFailed(String),
     #[error("fleet model download failed: {0}")]
diff --git a/crates/burn-central-fleet/src/session.rs b/crates/burn-central-fleet/src/session.rs
@@ -1,4 +1,7 @@
-use std::{path::PathBuf, sync::Arc};
+use std::{
+    path::PathBuf,
+    sync::{Arc, RwLock},
+};
 
 use burn_central_client::{Env, FleetClient};
 use directories::{BaseDirs, ProjectDirs};
@@ -15,6 +18,7 @@ pub struct FleetDeviceSession {
     registration_token: FleetRegistrationToken,
     identity_key: String,
     state: state::FleetState,
+    telemetry_auth_token: Arc<RwLock<Option<String>>>,
     client: FleetClient,
     fleet_key: String,
     store: state::FleetLocalStateStore,
@@ -40,18 +44,24 @@ impl FleetDeviceSession {
         let client = FleetClient::new(env.clone());
         let store = state::FleetLocalStateStore::new(root_dir.clone());
         let identity_key = store.load_or_create_machine_identity_key()?;
+        let state = store.load_fleet_state(&fleet_key)?.unwrap_or_default();
+        let telemetry_auth_token = Arc::new(RwLock::new(
+            state
+                .auth_token()
+                .filter(|auth| auth.is_valid())
+                .map(|auth| auth.token().to_string()),
+        ));
         let telemetry = TelemetryPipeline::get_or_init(
             fleet_key.clone(),
-            registration_token.clone(),
-            identity_key.clone(),
+            telemetry_auth_token.clone(),
             client.clone(),
             root_dir,
         )?;
-        let state = store.load_fleet_state(&fleet_key)?.unwrap_or_default();
         let fleet_device = FleetDeviceSession {
             registration_token,
             identity_key,
             state,
+            telemetry_auth_token,
             client,
             fleet_key,
             store,
@@ -98,18 +108,48 @@ impl FleetDeviceSession {
             ?metadata,
             "syncing fleet device with fleet management service"
         );
+
+        let should_refresh_token = match self.state.auth_token() {
+            Some(auth) if auth.is_valid() => {
+                tracing::debug!(
+                    "using existing auth token with ttl {} seconds",
+                    auth.expires_in_seconds().unwrap()
+                );
+                false
+            }
+            Some(_) => {
+                tracing::info!(
+                    "existing auth token expired, requesting a new one from fleet management service"
+                );
+                true
+            }
+            None => {
+                tracing::info!(
+                    "no existing auth token, requesting new one from fleet management service"
+                );
+                true
+            }
+        };
+
+        if should_refresh_token {
+            self.refresh_auth_token(metadata.clone())?;
+            tracing::info!("successfully refreshed auth token");
+        }
+        let auth_token = self
+            .state
+            .auth_token()
+            .ok_or_else(|| FleetError::SyncFailed("missing auth token after refresh".to_string()))?
+            .token()
+            .to_string();
+
         let snapshot = self
             .client
-            .sync(
-                self.registration_token.clone(),
-                self.identity_key.clone(),
-                metadata,
-            )
+            .sync(&auth_token, metadata)
             .map_err(|e| FleetError::SyncFailed(e.to_string()))?;
 
         let download = self
             .client
-            .model_download(self.registration_token.clone(), self.identity_key.clone())
+            .model_download(&auth_token)
             .map_err(|e| FleetError::DownloadFailed(e.to_string()))?;
 
         model::ensure_cached_model(
@@ -129,6 +169,28 @@ impl FleetDeviceSession {
             .save_fleet_state(&self.fleet_key, &self.state)
             .map_err(FleetError::from)
     }
+
+    fn refresh_auth_token(&mut self, metadata: Option<DeviceMetadata>) -> Result<(), FleetError> {
+        let auth_response = self
+            .client
+            .register(
+                self.registration_token.clone(),
+                self.identity_key.clone(),
+                metadata,
+            )
+            .map_err(|e| FleetError::RegistrationFailed(e.to_string()))?;
+
+        self.state
+            .set_auth_token(auth_response.access_token, auth_response.expires_in_seconds);
+
+        let mut telemetry_auth_token = self
+            .telemetry_auth_token
+            .write()
+            .map_err(|_| FleetError::SyncFailed("telemetry auth token lock poisoned".to_string()))?;
+        *telemetry_auth_token = self.state.auth_token().map(|auth| auth.token().to_string());
+
+        Ok(())
+    }
 }
 
 /// Get the default cache directory for a given environment.
diff --git a/crates/burn-central-fleet/src/state.rs b/crates/burn-central-fleet/src/state.rs
@@ -7,9 +7,48 @@ use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AuthToken {
+    /// The authentication token string to be used in requests to the fleet management service.
+    token: String,
+    /// The time-to-live of the authentication token, in seconds.
+    ttl_seconds: u64,
+    /// The timestamp of when the authentication token was last updated.
+    updated_at: String,
+}
+
+impl AuthToken {
+    /// Check if the authentication token is still valid based on its TTL and last updated timestamp.
+    pub fn is_valid(&self) -> bool {
+        let updated_at = match chrono::DateTime::parse_from_rfc3339(&self.updated_at) {
+            Ok(dt) => dt,
+            Err(_) => return false,
+        };
+
+        let expires_at = updated_at + chrono::Duration::seconds(self.ttl_seconds as i64);
+        let now = chrono::Utc::now();
+        now < expires_at
+    }
+
+    /// Calculate the number of seconds until the authentication token expires, or return None if the timestamp is invalid.
+    pub fn expires_in_seconds(&self) -> Option<i64> {
+        let updated_at = chrono::DateTime::parse_from_rfc3339(&self.updated_at).ok()?;
+        let expires_at = updated_at + chrono::Duration::seconds(self.ttl_seconds as i64);
+        let now = chrono::Utc::now().with_timezone(&updated_at.timezone());
+        Some((expires_at - now).num_seconds())
+    }
+
+    /// Get a reference to the authentication token string.
+    pub fn token(&self) -> &str {
+        &self.token
+    }
+}
+
 /// The state of a device in the fleet management system, as stored on the device and synced with the fleet management service.
 #[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct FleetState {
+    /// The current authentication token for communicating with the fleet management service, if any.
+    auth: Option<AuthToken>,
     /// The timestamp of the last update received by the device.
     updated_at: String,
     /// The id of the model version currently active on the device. Should be updated by the device when a new model version is activated.
@@ -21,6 +60,7 @@ pub struct FleetState {
 impl Default for FleetState {
     fn default() -> Self {
         Self {
+            auth: None,
             updated_at: chrono::Utc::now().to_rfc3339(),
             active_model_version_id: String::new(),
             runtime_config: serde_json::json!({}),
@@ -29,6 +69,15 @@ impl Default for FleetState {
 }
 
 impl FleetState {
+    pub fn set_auth_token(&mut self, token: String, ttl_seconds: u64) {
+        let now = chrono::Utc::now().to_rfc3339();
+        self.auth = Some(AuthToken {
+            token,
+            ttl_seconds,
+            updated_at: now,
+        });
+    }
+
     pub fn update(&mut self, model_version_id: String, runtime_config: serde_json::Value) {
         self.updated_at = chrono::Utc::now().to_rfc3339();
         self.active_model_version_id = model_version_id;
@@ -42,7 +91,12 @@ impl FleetState {
     pub fn runtime_config(&self) -> &serde_json::Value {
         &self.runtime_config
     }
+
+    pub fn auth_token(&self) -> Option<&AuthToken> {
+        self.auth.as_ref()
+    }
 }
+
 #[derive(Serialize, Deserialize, Clone, Debug)]
 struct IdentityState {
     identity_key: String,
@@ -186,6 +240,11 @@ mod tests {
 
         let fleet_key = fleet_key_from_registration_token("reg-token");
         let state = FleetState {
+            auth: Some(AuthToken {
+                token: "auth-token".to_string(),
+                ttl_seconds: 3600,
+                updated_at: "2026-01-01T00:00:00Z".to_string(),
+            }),
             updated_at: "2026-01-01T00:00:00Z".to_string(),
             active_model_version_id: "model-v1".to_string(),
             runtime_config: serde_json::json!({ "sample_rate": 0.2 }),
diff --git a/crates/burn-central-fleet/src/telemetry/pipeline/mod.rs b/crates/burn-central-fleet/src/telemetry/pipeline/mod.rs
@@ -6,6 +6,7 @@ use std::{
     sync::{
         Arc,
         atomic::{AtomicUsize, Ordering},
+        RwLock,
     },
     time::Duration,
 };
@@ -80,8 +81,7 @@ pub struct TelemetryPipeline {
 impl TelemetryPipeline {
     pub fn get_or_init(
         fleet_key: String,
-        registration_token: String,
-        identity_key: String,
+        auth_token: Arc<RwLock<Option<String>>>,
         client: FleetClient,
         root_dir: PathBuf,
     ) -> Result<Arc<Self>, TelemetryPipelineError> {
@@ -96,8 +96,7 @@ impl TelemetryPipeline {
         let recorder = global_recorder_handle();
         let pipeline = Arc::new(Self::start(
             fleet_key.clone(),
-            registration_token,
-            identity_key,
+            auth_token,
             client,
             recorder,
             root_dir,
@@ -112,8 +111,7 @@ impl TelemetryPipeline {
 
     fn start(
         fleet_key: String,
-        registration_token: String,
-        identity_key: String,
+        auth_token: Arc<RwLock<Option<String>>>,
         client: FleetClient,
         recorder: RecorderHandle,
         root_dir: PathBuf,
@@ -147,11 +145,7 @@ impl TelemetryPipeline {
 
         let shipper_handle = shipper::start(
             outbox,
-            Arc::new(shipper::BurnCentralFleetShipperTransport::new(
-                registration_token,
-                identity_key,
-                client,
-            )),
+            Arc::new(shipper::BurnCentralFleetShipperTransport::new(auth_token, client)),
             Duration::from_secs(5),
         );
 
diff --git a/crates/burn-central-fleet/src/telemetry/pipeline/shipper.rs b/crates/burn-central-fleet/src/telemetry/pipeline/shipper.rs
@@ -1,5 +1,5 @@
-use std::sync::Arc;
 use std::sync::mpsc::{RecvTimeoutError, Sender, channel};
+use std::sync::{Arc, RwLock};
 use std::time::Duration;
 
 use burn_central_client::FleetClient;
@@ -20,22 +20,13 @@ pub trait ShipperTransport: Send + Sync {
 }
 
 pub struct BurnCentralFleetShipperTransport {
-    registration_token: String,
-    identity_key: String,
+    auth_token: Arc<RwLock<Option<String>>>,
     client: FleetClient,
 }
 
 impl BurnCentralFleetShipperTransport {
-    pub fn new(
-        registration_token: impl Into<String>,
-        identity_key: impl Into<String>,
-        client: FleetClient,
-    ) -> Self {
-        Self {
-            registration_token: registration_token.into(),
-            identity_key: identity_key.into(),
-            client,
-        }
+    pub fn new(auth_token: Arc<RwLock<Option<String>>>, client: FleetClient) -> Self {
+        Self { auth_token, client }
     }
 }
 
@@ -124,9 +115,15 @@ impl ShipperTransport for BurnCentralFleetShipperTransport {
             metrics,
             logs,
         };
+        let auth_token = self
+            .auth_token
+            .read()
+            .map_err(|_| "telemetry auth token lock poisoned".to_string())?
+            .clone()
+            .ok_or_else(|| "missing auth token for telemetry ingestion".to_string())?;
 
         self.client
-            .ingest_telemetry(&self.registration_token, &self.identity_key, telemetry)
+            .ingest_telemetry(auth_token, telemetry)
             .map_err(|e| format!("failed to send telemetry events to Burn Central Fleet: {e}"))
     }
 }
