fix: retry transient OIDC refresh failures instead of disabling the controller

crates/agent/src/daemon.rs

@@ -379,25 +379,18 @@ where
         let state_dir = args.state_dir.clone();
         let oidc_callback_listen = args.oidc_callback_listen;
         let remote_enrollment = enrollment.clone();
-        tokio::spawn(async move {
-            if let Err(error) = remote::run(
-                controller,
-                remote_discovery,
-                state_dir,
-                oidc_callback_listen,
-                reconciler,
-                remote_enrollment.clone(),
-                remote::Requests {
-                    telemetry: telemetry_receiver,
-                    logout: logout_receiver,
-                },
-            )
-            .await
-            {
-                remote_enrollment.set("failed").await;
-                tracing::error!(error = %format!("{error:#}"), "controller integration disabled");
-            }
-        });
+        tokio::spawn(remote::run(
+            controller,
+            remote_discovery,
+            state_dir,
+            oidc_callback_listen,
+            reconciler,
+            remote_enrollment,
+            remote::Requests {
+                telemetry: telemetry_receiver,
+                logout: logout_receiver,
+            },
+        ));
     }
     let app = api::router(api::AppState {
         config,

crates/agent/src/remote.rs

@@ -1,4 +1,5 @@
 use std::{
+    convert::Infallible,
     future::Future,
     net::SocketAddr,
     path::{Path, PathBuf},
@@ -62,11 +63,77 @@ pub async fn run(
     reconciler: Reconciler,
     enrollment: EnrollmentState,
     requests: Requests,
-) -> anyhow::Result<()> {
+) {
     let Requests {
         mut telemetry,
         mut logout,
     } = requests;
+    let mut delay = INITIAL_RETRY_DELAY;
+    loop {
+        let started = time::Instant::now();
+        let Err(error) = run_session(
+            &controller,
+            &mut discovered,
+            &state_dir,
+            oidc_callback_listen,
+            &reconciler,
+            &enrollment,
+            &mut telemetry,
+            &mut logout,
+        )
+        .await;
+        // A session that stayed up for a while failed for a new reason; do not
+        // penalize it with the backoff accumulated by earlier failures.
+        if started.elapsed() > MAX_RETRY_DELAY {
+            delay = INITIAL_RETRY_DELAY;
+        }
+        enrollment.set("failed").await;
+        tracing::error!(
+            controller = %controller.address,
+            retry_in_seconds = delay.as_secs(),
+            error = %format!("{error:#}"),
+            "controller integration failed; restarting"
+        );
+        tokio::select! {
+            _ = time::sleep(delay) => delay = next_retry_delay(delay),
+            Some(request) = logout.recv() => {
+                logout_between_sessions(request, &state_dir.join("identity.json"), &enrollment)
+                    .await;
+                delay = INITIAL_RETRY_DELAY;
+            }
+        }
+    }
+}
+
+async fn logout_between_sessions(
+    request: LogoutRequest,
+    identity_path: &Path,
+    enrollment: &EnrollmentState,
+) {
+    match identity::load(identity_path) {
+        Ok(Some(identity)) => {
+            complete_logout(request, identity_path, &identity, enrollment).await;
+        }
+        Ok(None) => complete_unenrolled_logout(request, enrollment).await,
+        Err(error) => {
+            let _ = request
+                .completion
+                .send(Err(format!("read local organization identity: {error:#}")));
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+async fn run_session(
+    controller: &ControllerConnectionConfig,
+    discovered: &mut watch::Receiver<Arc<AgentDiscovery>>,
+    state_dir: &Path,
+    oidc_callback_listen: Option<SocketAddr>,
+    reconciler: &Reconciler,
+    enrollment: &EnrollmentState,
+    telemetry: &mut mpsc::Receiver<ModelTelemetryEvent>,
+    logout: &mut mpsc::Receiver<LogoutRequest>,
+) -> anyhow::Result<Infallible> {
     let identity_path = state_dir.join("identity.json");
     loop {
         let mut identity = match identity::load(&identity_path)? {
@@ -77,10 +144,10 @@ pub async fn run(
             None => {
                 let identity = enroll_with_retry(
                     &controller.address,
-                    &enrollment,
-                    &mut logout,
+                    enrollment,
+                    logout,
                     INITIAL_RETRY_DELAY,
-                    || oidc::enroll(&controller, &enrollment, oidc_callback_listen),
+                    || oidc::enroll(controller, enrollment, oidc_callback_listen),
                 )
                 .await;
                 identity::save(&identity_path, &identity)?;
@@ -95,7 +162,7 @@ pub async fn run(
             let refresh_result = tokio::select! {
                 result = refresh_oauth_if_needed(&mut identity, &identity_path) => result,
                 Some(request) = logout.recv() => {
-                    if complete_logout(request, &identity_path, &identity, &enrollment).await {
+                    if complete_logout(request, &identity_path, &identity, enrollment).await {
                         break;
                     }
                     continue;
@@ -115,10 +182,29 @@ pub async fn run(
                     );
                     break;
                 }
+                Err(error) if is_transient_refresh_failure(&error) => {
+                    // Network failures and token endpoint outages are transient: the
+                    // stored refresh token is still valid, so keep the identity and retry
+                    // instead of tearing down the controller integration.
+                    warn!(
+                        controller = %controller.address,
+                        retry_in_seconds = delay.as_secs(),
+                        error = %format!("{error:#}"),
+                        "OIDC access token refresh failed; retaining current identity and retrying"
+                    );
+                    if wait_before_retry(delay, logout, &identity_path, &identity, enrollment).await
+                    {
+                        break;
+                    }
+                    delay = next_retry_delay(delay);
+                    continue;
+                }
+                // Local identity store failures and unsupported token types are not
+                // fixed by waiting; let the supervisor report them as failed.
                 Err(error) => return Err(error),
             }
             if certificate_needs_renewal(&identity) {
-                match renew_device_certificate(&controller, &identity).await {
+                match renew_device_certificate(controller, &identity).await {
                     Ok(renewed) => {
                         identity::save(&identity_path, &renewed)?;
                         identity = renewed;
@@ -132,15 +218,15 @@ pub async fn run(
             identity::save(&identity_path, &identity)?;
             let connection = tokio::select! {
                 result = connect(
-                    &controller,
+                    controller,
                     &identity,
-                    &mut discovered,
-                    &state_dir,
-                    &reconciler,
-                    &mut telemetry,
+                    discovered,
+                    state_dir,
+                    reconciler,
+                    telemetry,
                 ) => Some(result),
                 Some(request) = logout.recv() => {
-                    if complete_logout(request, &identity_path, &identity, &enrollment).await {
+                    if complete_logout(request, &identity_path, &identity, enrollment).await {
                         None
                     } else {
                         continue;
@@ -175,20 +261,31 @@ pub async fn run(
                 }
             }
 
-            let logged_out = tokio::select! {
-                _ = time::sleep(delay) => false,
-                Some(request) = logout.recv() => {
-                    complete_logout(request, &identity_path, &identity, &enrollment).await
-                }
-            };
-            if logged_out {
+            if wait_before_retry(delay, logout, &identity_path, &identity, enrollment).await {
                 break;
             }
             delay = next_retry_delay(delay);
         }
     }
 }
 
+/// Sleeps for `delay`, completing any logout request that arrives meanwhile.
+/// Returns `true` when the local session was logged out.
+async fn wait_before_retry(
+    delay: Duration,
+    logout: &mut mpsc::Receiver<LogoutRequest>,
+    identity_path: &Path,
+    identity: &Identity,
+    enrollment: &EnrollmentState,
+) -> bool {
+    tokio::select! {
+        _ = time::sleep(delay) => false,
+        Some(request) = logout.recv() => {
+            complete_logout(request, identity_path, identity, enrollment).await
+        }
+    }
+}
+
 async fn enroll_with_retry<F, Fut, T>(
     controller: &str,
     enrollment: &EnrollmentState,
@@ -278,6 +375,14 @@ fn is_oauth_refresh_rejected(error: &anyhow::Error) -> bool {
         })
 }
 
+/// Refresh failures that came from talking to the token endpoint (connection
+/// errors, timeouts, non-rejecting HTTP statuses) and may succeed on retry.
+fn is_transient_refresh_failure(error: &anyhow::Error) -> bool {
+    error
+        .chain()
+        .any(|cause| cause.downcast_ref::<reqwest::Error>().is_some())
+}
+
 pub async fn llm_gateway_credential(
     controller: &ControllerConnectionConfig,
     state_dir: &Path,
@@ -700,8 +805,9 @@ mod tests {
     use tokio::sync::mpsc;
 
     use super::{
-        MAX_RETRY_DELAY, enroll_with_retry, is_oauth_refresh_rejected, is_unauthenticated,
-        next_inventory, next_retry_delay, normalize_hostname,
+        MAX_RETRY_DELAY, enroll_with_retry, is_oauth_refresh_rejected,
+        is_transient_refresh_failure, is_unauthenticated, next_inventory, next_retry_delay,
+        normalize_hostname,
     };
     use crate::enrollment::EnrollmentState;
     use agentdesktop_core::model::Discovery as AgentDiscovery;
@@ -775,6 +881,57 @@ mod tests {
         assert!(!is_oauth_refresh_rejected(&anyhow::anyhow!(
             "network unavailable"
         )));
+        assert!(is_transient_refresh_failure(
+            &status_error(reqwest::StatusCode::SERVICE_UNAVAILABLE)
+                .context("OIDC token endpoint rejected refresh token")
+        ));
+    }
+
+    /// Identity store failures are not retried in place; they must reach the
+    /// supervisor so enrollment is reported as failed.
+    #[test]
+    fn local_identity_failures_are_not_transient() {
+        let error = anyhow::Error::new(std::io::Error::from(std::io::ErrorKind::PermissionDenied))
+            .context("persist rotated OIDC refresh token");
+        assert!(!is_transient_refresh_failure(&error));
+        assert!(!is_transient_refresh_failure(&anyhow::anyhow!(
+            "OIDC token endpoint returned unsupported token type"
+        )));
+    }
+
+    #[tokio::test]
+    async fn network_failures_during_refresh_are_transient() {
+        let closed = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+        let closed_address = closed.local_addr().unwrap();
+        drop(closed);
+        let connect_error: anyhow::Error = reqwest::Client::new()
+            .post(format!("http://{closed_address}/token"))
+            .send()
+            .await
+            .unwrap_err()
+            .into();
+        let connect_error = connect_error.context("refresh OIDC access token");
+        assert!(!is_oauth_refresh_rejected(&connect_error));
+        assert!(is_transient_refresh_failure(&connect_error));
+
+        let silent = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let silent_address = silent.local_addr().unwrap();
+        let _accepted = tokio::spawn(async move {
+            let _connection = silent.accept().await;
+            std::future::pending::<()>().await;
+        });
+        let timeout_error: anyhow::Error = reqwest::Client::builder()
+            .timeout(Duration::from_millis(50))
+            .build()
+            .unwrap()
+            .post(format!("http://{silent_address}/token"))
+            .send()
+            .await
+            .unwrap_err()
+            .into();
+        let timeout_error = timeout_error.context("refresh OIDC access token");
+        assert!(!is_oauth_refresh_rejected(&timeout_error));
+        assert!(is_transient_refresh_failure(&timeout_error));
     }
 
     #[test]