fix: defer OIDC token acquisition until after Nix evaluation

[flexiondotorg]

For large flakes, Nix evaluation (nix flake show --all-systems + output path enumeration) can take 8+ minutes, but GitHub Actions OIDC tokens expire after ~5 minutes. Previously, the token was obtained in PushContext::from_cli_and_env() before evaluation began, causing 401 Unauthorized errors when the token was finally used for API calls.

Move token acquisition to just before FlakeHubClient::new() in execute(), after all Nix evaluation completes. Introduce a TokenContext enum that captures environment-specific data needed for each auth method (GitHub, GitLab, Generic, LocalGitHub) without immediately fetching the token. The new acquire_auth_token() method on PushContext handles deferred retrieval.

Fixes #275

Summary by CodeRabbit

New Features

• Enhanced authentication handling – Authentication tokens are now refreshed before the first API call, improving security and reliability.
• Improved multi-platform support – Better token acquisition for GitHub, GitLab, generic runners, and local environments with environment-specific handling.

[coderabbitai]

📝 Walkthrough

Walkthrough

Moved OIDC token acquisition from early PushContext construction to immediately after flake evaluation completes. Introduced a TokenContext enum that defers token retrieval by environment type (GitHub, GitLab, Generic, LocalGitHub), then acquires fresh tokens via a new async method just before FlakeHubClient initialization to prevent token expiry during long evaluations.

Changes

Cohort / File(s)	Summary
Main Entry Point src/main.rs	Added call to acquire_auth_token() after constructing PushContext and before initializing FlakeHubClient to fetch a fresh token post-evaluation.
Token Acquisition Logic src/push_context.rs	Introduced TokenContext enum to encapsulate token acquisition strategy by environment (GitHub, GitLab, Generic, LocalGitHub). Replaced auth_token: String field with token_context: TokenContext. Added async acquire_auth_token() method that defers token fetching until after flake evaluation, including special handling for LocalGitHub to reset context after token generation.

Sequence Diagram

sequenceDiagram
    participant Main as Main (execute)
    participant PushCtx as PushContext
    participant TokenMgr as Token Acquisition
    participant FlakeEval as Flake Evaluation
    participant Client as FlakeHubClient
    participant API as FlakeHub API

    Main->>PushCtx: from_cli_and_env()<br/>(stores TokenContext)
    Main->>FlakeEval: Run nix evaluation<br/>(can take 8+ minutes)
    FlakeEval-->>Main: Evaluation complete
    Main->>TokenMgr: acquire_auth_token()
    
    alt GitHub Actions
        TokenMgr->>TokenMgr: get_actions_id_bearer_token()
    else GitLab Runner
        TokenMgr->>TokenMgr: get_runner_bearer_token()
    else Generic
        TokenMgr->>TokenMgr: Read FLAKEHUB_PUSH_OIDC_TOKEN env
    else LocalGitHub
        TokenMgr->>TokenMgr: Generate fake token
        TokenMgr->>TokenMgr: Reset to Generic
    end
    
    TokenMgr-->>Main: (fresh_token, ctx)
    Main->>Client: FlakeHubClient::new(fresh_token)
    Client->>API: /token/status (with fresh token)
    API-->>Client: 200 OK

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Poem

🐰 A token that traveled through time, now refreshed and sublime!
No more 401 heartache—evaluation's done, and we grab it just right.
Fresh OIDC upon the api call's plight! 🥕✨

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: deferring OIDC token acquisition until after Nix evaluation completes, which directly addresses the core problem.
Linked Issues check	✅ Passed	The PR successfully implements the suggested fix from issue #275: moving token acquisition from PushContext::from_cli_and_env() to after evaluation via the new acquire_auth_token() method called before FlakeHubClient construction.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to deferring token acquisition: introducing TokenContext enum, adding acquire_auth_token() method, and updating token handling logic—no unrelated modifications detected.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Merge Conflict Detection	✅ Passed	✅ No merge conflicts detected when merging into main

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

No actionable comments were generated in the recent review. 🎉

🧹 Recent nitpick comments

src/push_context.rs (2)

195-257: Consider taking &mut self to avoid the verbose destructure-reconstruct pattern.

The full destructure on lines 199–207 and reconstruct on lines 246–254 exist solely because self is consumed by value. Since only the LocalGitHub arm actually needs to move data out of token_context, you could take &mut self and use std::mem::replace for that case, returning just the String token. This simplifies the caller too (let auth_token = ctx.acquire_auth_token().await?).

♻️ Sketch using &mut self

-    pub async fn acquire_auth_token(self) -> Result<(String, Self)> {
-        let PushContext {
-            flakehub_host,
-            token_context,
-            upload_name,
-            release_version,
-            error_if_release_conflicts,
-            metadata,
-            tarball,
-        } = self;
-
-        let (token, token_context) = match token_context {
-            TokenContext::GitHub { ref host } => {
+    pub async fn acquire_auth_token(&mut self) -> Result<String> {
+        let token = match &self.token_context {
+            TokenContext::GitHub { host } => {
                 let t = crate::github::get_actions_id_bearer_token(host)
                     .await
                     .wrap_err("Getting upload bearer token from GitHub")?;
-                (t, token_context)
+                t
             }
             TokenContext::GitLab => {
                 let t = crate::gitlab::get_runner_bearer_token()
                     .await
                     .wrap_err("Getting upload bearer token from GitLab")?;
-                (t, TokenContext::GitLab)
+                t
             }
             TokenContext::Generic => {
                 let t = std::env::var("FLAKEHUB_PUSH_OIDC_TOKEN")
                     .with_context(|| "missing FLAKEHUB_PUSH_OIDC_TOKEN environment variable")?;
-                (t, TokenContext::Generic)
+                t
             }
-            TokenContext::LocalGitHub {
-                jwt_issuer_uri,
-                project_owner,
-                repository,
-                github_graphql_data_result,
-            } => {
-                let t = flakehub_auth_fake::get_fake_bearer_token(
-                    &jwt_issuer_uri,
-                    &project_owner,
-                    &repository,
-                    github_graphql_data_result,
-                )
-                .await?;
-                (t, TokenContext::Generic)
+            TokenContext::LocalGitHub { .. } => {
+                let old = std::mem::replace(&mut self.token_context, TokenContext::Generic);
+                let TokenContext::LocalGitHub {
+                    jwt_issuer_uri,
+                    project_owner,
+                    repository,
+                    github_graphql_data_result,
+                } = old else { unreachable!() };
+                flakehub_auth_fake::get_fake_bearer_token(
+                    &jwt_issuer_uri,
+                    &project_owner,
+                    &repository,
+                    github_graphql_data_result,
+                )
+                .await?
             }
         };
-
-        let ctx = PushContext {
-            flakehub_host,
-            token_context,
-            upload_name,
-            release_version,
-            error_if_release_conflicts,
-            metadata,
-            tarball,
-        };
-
-        Ok((token, ctx))
+        Ok(token)
     }

---

110-112: host in TokenContext::GitHub duplicates flakehub_host — consider referencing the existing field.

TokenContext::GitHub { host } stores a clone of cli.host (line 111), which is the same value stored in PushContext::flakehub_host (line 175). If you switch acquire_auth_token to &mut self as suggested above, the GitHub arm could simply reference self.flakehub_host directly, removing the need for the host field in TokenContext::GitHub entirely and making it a unit variant like GitLab/Generic.

Tip

Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}