Experimental workspace plumbing

Signed-off-by: William Woodruff <william@astral.sh>

Author: woodruffw

crates/uv-auth/src/index.rs

@@ -107,6 +107,11 @@ impl Indexes {
             .unwrap_or(AuthPolicy::Auto)
     }
 
+    /// Iterate over all indexes.
+    pub fn iter(&self) -> impl Iterator<Item = &Index> {
+        self.0.iter()
+    }
+
     fn find_prefix_index(&self, url: &Url) -> Option<&Index> {
         self.0.iter().find(|&index| index.is_prefix_for(url))
     }

crates/uv-auth/src/middleware.rs

@@ -5,6 +5,7 @@ use http::{Extensions, StatusCode};
 use netrc::Netrc;
 use reqwest::{Request, Response};
 use reqwest_middleware::{ClientWithMiddleware, Error, Middleware, Next};
+use rustc_hash::FxHashMap;
 use tokio::sync::Mutex;
 use tracing::{debug, trace, warn};
 
@@ -120,15 +121,6 @@ impl TextStoreMode {
     }
 }
 
-#[derive(Debug, Clone)]
-enum TokenState {
-    /// The token state has not yet been initialized from the store.
-    Uninitialized,
-    /// The token state has been initialized, and the store either returned tokens or `None` if
-    /// the user has not yet authenticated.
-    Initialized(Option<AccessToken>),
-}
-
 #[derive(Clone)]
 enum S3CredentialState {
     /// The S3 credential state has not yet been initialized.
@@ -166,8 +158,13 @@ pub struct AuthMiddleware {
     base_client: Option<ClientWithMiddleware>,
     /// The pyx token store to use for persistent credentials.
     pyx_token_store: Option<PyxTokenStore>,
-    /// Tokens to use for persistent credentials.
-    pyx_token_state: Mutex<TokenState>,
+    /// Per-workspace token state.
+    ///
+    /// The key is the workspace name (e.g., `"acme"`), or `None` for CDN/non-workspace URLs.
+    /// Each workspace gets its own entry so that Trusted Access tokens — which are minted
+    /// per-workspace — are cached and refreshed independently. A missing entry means the token
+    /// has not yet been fetched for that workspace.
+    pyx_token_state: Mutex<FxHashMap<Option<String>, Option<AccessToken>>>,
     /// Cached S3 credentials to avoid running the credential helper multiple times.
     s3_credential_state: Mutex<S3CredentialState>,
     /// Cached GCS credentials to avoid running the credential helper multiple times.
@@ -193,7 +190,7 @@ impl AuthMiddleware {
             only_authenticated: false,
             base_client: None,
             pyx_token_store: None,
-            pyx_token_state: Mutex::new(TokenState::Uninitialized),
+            pyx_token_state: Mutex::new(FxHashMap::default()),
             s3_credential_state: Mutex::new(S3CredentialState::Uninitialized),
             gcs_credential_state: Mutex::new(GcsCredentialState::Uninitialized),
             preview: Preview::default(),
@@ -752,14 +749,37 @@ impl AuthMiddleware {
         if let Some(base_client) = self.base_client.as_ref() {
             if let Some(token_store) = self.pyx_token_store.as_ref() {
                 if token_store.is_known_url(url) {
-                    let mut token_state = self.pyx_token_state.lock().await;
+                    // Derive the workspace from the URL (e.g., `acme` from
+                    // `https://api.pyx.dev/simple/acme/main/`). CDN URLs have no workspace in
+                    // their path, so this returns `None` for them.
+                    let workspace = token_store.workspace_for_url(url).map(str::to_owned);
+
+                    let mut token_state_map = self.pyx_token_state.lock().await;
 
-                    // If the token store is uninitialized, initialize it.
-                    let token = match *token_state {
-                        TokenState::Uninitialized => {
+                    let token = if let Some(token) = token_state_map.get(&workspace) {
+                        token.clone()
+                    } else {
+                        // For CDN URLs (workspace=None), try to reuse any already-initialized
+                        // workspace token rather than bootstrapping from scratch. Any valid
+                        // pyx token should authorize CDN downloads.
+                        let existing = if workspace.is_none() {
+                            token_state_map.values().find_map(std::clone::Clone::clone)
+                        } else {
+                            None
+                        };
+
+                        if let Some(token) = existing {
+                            // Cache under the None key for subsequent CDN requests.
+                            token_state_map.insert(None, Some(token.clone()));
+                            Some(token)
+                        } else {
                             trace!("Initializing token store for {url}");
                             let generated = match token_store
-                                .access_token(base_client, DEFAULT_TOLERANCE_SECS)
+                                .access_token(
+                                    base_client,
+                                    DEFAULT_TOLERANCE_SECS,
+                                    workspace.as_deref(),
+                                )
                                 .await
                             {
                                 Ok(Some(token)) => Some(token),
@@ -769,10 +789,9 @@ impl AuthMiddleware {
                                     None
                                 }
                             };
-                            *token_state = TokenState::Initialized(generated.clone());
+                            token_state_map.insert(workspace.clone(), generated.clone());
                             generated
                         }
-                        TokenState::Initialized(ref tokens) => tokens.clone(),
                     };
 
                     let credentials = token.map(|token| {

crates/uv-auth/src/pyx.rs

@@ -253,6 +253,15 @@ impl PyxTokenStore {
         })
     }
 
+    /// Extract the workspace name from a pyx Simple API URL.
+    ///
+    /// Pyx Simple API index URLs have the form `{api}/simple/{workspace}/{view}` (e.g.,
+    /// `https://api.pyx.dev/simple/acme/main`). Returns `None` if the URL does not match this
+    /// store's API URL or does not have the expected path shape.
+    pub fn workspace_for_url<'a>(&self, url: &'a DisplaySafeUrl) -> Option<&'a str> {
+        workspace_from_simple_url(url, &self.api)
+    }
+
     /// Return the root directory for the token store.
     pub fn root(&self) -> &Path {
         &self.root
@@ -271,18 +280,22 @@ impl PyxTokenStore {
     ///
     /// If no access token is found, but an API key is present, the API key will be used to
     /// bootstrap an access token.
+    ///
+    /// `workspace` is the pyx workspace name, required to bootstrap a Trusted Access token. If
+    /// `None`, Trusted Access bootstrapping is skipped.
     pub async fn access_token(
         &self,
         client: &ClientWithMiddleware,
         tolerance_secs: u64,
+        workspace: Option<&str>,
     ) -> Result<Option<AccessToken>, TokenStoreError> {
         // If the access token is already set in the environment, return it.
         if let Some(access_token) = read_pyx_auth_token() {
             return Ok(Some(access_token));
         }
 
         // Initialize the tokens from the store.
-        let tokens = self.init(client, tolerance_secs).await?;
+        let tokens = self.init(client, tolerance_secs, workspace).await?;
 
         // Extract the access token from the OAuth tokens or API key.
         Ok(tokens.map(AccessToken::from))
@@ -294,20 +307,26 @@ impl PyxTokenStore {
     ///
     /// If no access token is found, but an API key is present, the API key will be used to
     /// bootstrap an access token.
+    ///
+    /// `workspace` is the pyx workspace name, required to bootstrap a Trusted Access token. If
+    /// `None`, Trusted Access bootstrapping is skipped.
     pub async fn init(
         &self,
         client: &ClientWithMiddleware,
         tolerance_secs: u64,
+        workspace: Option<&str>,
     ) -> Result<Option<PyxTokens>, TokenStoreError> {
         match self.read().await? {
             Some(tokens) => {
                 // Refresh the tokens if they are expired.
-                let tokens = self.refresh(tokens, client, tolerance_secs).await?;
+                let tokens = self
+                    .refresh(tokens, client, tolerance_secs, workspace)
+                    .await?;
                 Ok(Some(tokens))
             }
             None => {
-                // If no tokens are present, bootstrap them from an API key.
-                self.bootstrap(client).await
+                // If no tokens are present, bootstrap them from an API key or Trusted Access.
+                self.bootstrap(client, workspace).await
             }
         }
     }
@@ -412,23 +431,31 @@ impl PyxTokenStore {
     }
 
     /// Bootstrap the tokens from an API key or from Trusted Access.
+    ///
+    /// `workspace` is the pyx workspace name, required to bootstrap a Trusted Access token. If
+    /// `None`, Trusted Access bootstrapping is skipped.
     async fn bootstrap(
         &self,
         client: &ClientWithMiddleware,
+        workspace: Option<&str>,
     ) -> Result<Option<PyxTokens>, TokenStoreError> {
         if let Some(api_key) = read_pyx_api_key() {
             // If an API key is present, use it to bootstrap the tokens.
             self.bootstrap_from_api_key(api_key, client).await.map(Some)
         } else {
             // Otherwise, attempt to bootstrap from Trusted Access.
-            self.bootstrap_from_trusted_access(client).await
+            self.bootstrap_from_trusted_access(client, workspace).await
         }
     }
 
     /// Bootstrap a [`PyxTokens`] from Trusted Access.
+    ///
+    /// `workspace` is the pyx workspace name. If `None`, Trusted Access bootstrapping is skipped,
+    /// since the workspace is required to construct the token endpoint.
     async fn bootstrap_from_trusted_access(
         &self,
         client: &ClientWithMiddleware,
+        workspace: Option<&str>,
     ) -> Result<Option<PyxTokens>, TokenStoreError> {
         #[derive(Debug, Clone, serde::Serialize)]
         struct RequestBody<'a> {
@@ -440,6 +467,12 @@ impl PyxTokenStore {
             token: AccessToken,
         }
 
+        // Trusted Access requires the workspace name to construct the token endpoint.
+        let Some(workspace) = workspace else {
+            debug!("Skipping Trusted Access: pyx workspace name is not known");
+            return Ok(None);
+        };
+
         // Get an OIDC token from the ambient environment (e.g., GitHub Actions).
         let detector = ambient_id::Detector::new_with_client(client.clone());
         let Some(id_token) = detector.detect("pyx:trusted-access").await? else {
@@ -449,7 +482,7 @@ impl PyxTokenStore {
 
         // Exchange the OIDC token for a Trusted Access token from pyx.
         let mut url = self.api.clone();
-        url.set_path("v1/trusted-access/TODO/mint-token");
+        url.set_path(&format!("v1/trusted-access/{workspace}/mint-token"));
 
         let request = client
             .request(reqwest::Method::POST, Url::from(url))
@@ -509,6 +542,7 @@ impl PyxTokenStore {
         tokens: PyxTokens,
         client: &ClientWithMiddleware,
         tolerance_secs: u64,
+        workspace: Option<&str>,
     ) -> Result<PyxTokens, TokenStoreError> {
         let reason = match tokens.check_fresh(tolerance_secs) {
             Ok(exp) => {
@@ -589,7 +623,7 @@ impl PyxTokenStore {
             }
             PyxTokens::TrustedAccess(_) => {
                 // Refreshing a Trusted Access token is the same as bootstrapping it.
-                self.bootstrap_from_trusted_access(client)
+                self.bootstrap_from_trusted_access(client, workspace)
                     .await?
                     .ok_or_else(|| {
                         // This can only happen if we're in an environment that previously
@@ -729,6 +763,30 @@ fn is_known_domain(url: &Url, api: &DisplaySafeUrl, cdn: &str) -> bool {
     is_known_url(url, api, cdn)
 }
 
+/// Extract the workspace name from a pyx Simple API URL.
+///
+/// Pyx Simple API URLs have the form `{api}/simple/{workspace}/{view}` (e.g.,
+/// `https://api.pyx.dev/simple/acme/main`). Returns the workspace segment when
+/// the URL matches the given `api` base URL.
+fn workspace_from_simple_url<'a>(url: &'a DisplaySafeUrl, api: &DisplaySafeUrl) -> Option<&'a str> {
+    // The URL must be on the same host/port as the API.
+    if url.scheme() != api.scheme()
+        || url.host_str() != api.host_str()
+        || url.port_or_known_default() != api.port_or_known_default()
+    {
+        return None;
+    }
+    // Path must be `/simple/{workspace}/{view}[/]`.
+    let mut segments = url.path_segments()?;
+    if segments.next()? != "simple" {
+        return None;
+    }
+    let workspace = segments.next().filter(|s| !s.is_empty())?;
+    // There must be at least a view segment after the workspace.
+    segments.next().filter(|s| !s.is_empty())?;
+    Some(workspace)
+}
+
 /// Returns `true` if the URL is on the default pyx domain.
 ///
 /// This is used in auth commands to recognize `pyx.dev` as a pyx domain even when
@@ -943,4 +1001,73 @@ mod tests {
             &Url::parse("https://beta.pyx.dev").unwrap()
         ));
     }
+
+    #[test]
+    fn test_workspace_from_simple_url() {
+        let api = DisplaySafeUrl::parse("https://api.pyx.dev").unwrap();
+
+        // Standard pyx simple index URL.
+        assert_eq!(
+            workspace_from_simple_url(
+                &DisplaySafeUrl::parse("https://api.pyx.dev/simple/acme/main").unwrap(),
+                &api
+            ),
+            Some("acme")
+        );
+
+        // Trailing slash is fine.
+        assert_eq!(
+            workspace_from_simple_url(
+                &DisplaySafeUrl::parse("https://api.pyx.dev/simple/acme/main/").unwrap(),
+                &api
+            ),
+            Some("acme")
+        );
+
+        // Must have a view segment after the workspace (bare /simple/acme is not a full index URL).
+        assert_eq!(
+            workspace_from_simple_url(
+                &DisplaySafeUrl::parse("https://api.pyx.dev/simple/acme").unwrap(),
+                &api
+            ),
+            None
+        );
+
+        // Non-simple path returns None.
+        assert_eq!(
+            workspace_from_simple_url(
+                &DisplaySafeUrl::parse("https://api.pyx.dev/v1/upload/acme/main").unwrap(),
+                &api
+            ),
+            None
+        );
+
+        // Different host returns None.
+        assert_eq!(
+            workspace_from_simple_url(
+                &DisplaySafeUrl::parse("https://other.pyx.dev/simple/acme/main").unwrap(),
+                &api
+            ),
+            None
+        );
+
+        // Different scheme returns None.
+        assert_eq!(
+            workspace_from_simple_url(
+                &DisplaySafeUrl::parse("http://api.pyx.dev/simple/acme/main").unwrap(),
+                &api
+            ),
+            None
+        );
+
+        // Custom API URL works the same way.
+        let custom_api = DisplaySafeUrl::parse("https://staging.example.com").unwrap();
+        assert_eq!(
+            workspace_from_simple_url(
+                &DisplaySafeUrl::parse("https://staging.example.com/simple/myorg/prod").unwrap(),
+                &custom_api
+            ),
+            Some("myorg")
+        );
+    }
 }

crates/uv/src/commands/auth/helper.rs

@@ -94,6 +94,7 @@ async fn credentials_for_url(
             .access_token(
                 client.for_host(pyx_store.api()).raw_client(),
                 DEFAULT_TOLERANCE_SECS,
+                pyx_store.workspace_for_url(url),
             )
             .await
             .context("Authentication failure")?

crates/uv/src/commands/auth/token.rs

@@ -91,7 +91,7 @@ async fn pyx_refresh(store: &PyxTokenStore, client: &BaseClient, printer: Printe
     // Retrieve the token store.
     // Use zero tolerance to force a refresh.
     let token = match store
-        .access_token(client.for_host(store.api()).raw_client(), 0)
+        .access_token(client.for_host(store.api()).raw_client(), 0, None)
         .await
     {
         // If the tokens were successfully refreshed, return them.