Support Trusted Publishing with pyx

.github/workflows/ci.yml

@@ -304,7 +304,7 @@ jobs:
           repository-url: "https://test.pypi.org/legacy/"
           packages-dir: "astral-test-pypa-gh-action/dist"
 
-      - name: "Request GitLab OIDC token for impersonation"
+      - name: "Request GitLab OIDC tokens for impersonation"
         uses: digital-blueprint/gitlab-pipeline-trigger-action@20e77989b24af658ba138a0aa5291bdc657f1505 # v1.3.0
         with:
           host: gitlab.com
@@ -316,21 +316,31 @@ jobs:
           fail_if_no_artifacts: true
           download_path: ./gitlab-artifacts
 
-      - name: "Load GitLab OIDC token from GitLab job artifacts"
+      - name: "Load GitLab OIDC tokens from GitLab job artifacts"
         id: load-gitlab-oidc-token
         run: |
-          # we expect ./gitlab-artifacts/*/artifacts/id-token to exist
-          id_token_file=$(find ./gitlab-artifacts -type f -name id-token | head -n 1)
-          if [ -z "${id_token_file}" ]; then
-            echo "No id-token file found in GitLab artifacts"
+          # we expect ./gitlab-artifacts/*/artifacts/pypi-id-token to exist
+          pypi_id_token_file=$(find ./gitlab-artifacts -type f -name pypi-id-token | head -n 1)
+          if [ -z "${pypi_id_token_file}" ]; then
+            echo "No pypi-id-token file found in GitLab artifacts"
             exit 1
           fi
-          GITLAB_OIDC_TOKEN=$(cat "${id_token_file}")
+          GITLAB_PYPI_OIDC_TOKEN=$(cat "${pypi_id_token_file}")
 
-          # Add a secret mask for the token.
-          echo "::add-mask::$GITLAB_OIDC_TOKEN"
+          # we expect ./gitlab-artifacts/*/artifacts/pyx-id-token to exist
+          pyx_id_token_file=$(find ./gitlab-artifacts -type f -name pyx-id-token | head -n 1)
+          if [ -z "${pyx_id_token_file}" ]; then
+            echo "No pyx-id-token file found in GitLab artifacts"
+            exit 1
+          fi
+          GITLAB_PYX_OIDC_TOKEN=$(cat "${pyx_id_token_file}")
+
+          # Add secret masks for the tokens.
+          echo "::add-mask::$GITLAB_PYPI_OIDC_TOKEN"
+          echo "::add-mask::$GITLAB_PYX_OIDC_TOKEN"
 
-          echo "GITLAB_OIDC_TOKEN=${GITLAB_OIDC_TOKEN}" >> "${GITHUB_OUTPUT}"
+          echo "GITLAB_PYPI_OIDC_TOKEN=${GITLAB_PYPI_OIDC_TOKEN}" >> "${GITHUB_OUTPUT}"
+          echo "GITLAB_PYX_OIDC_TOKEN=${GITLAB_PYX_OIDC_TOKEN}" >> "${GITHUB_OUTPUT}"
 
       - name: "Add password to keyring"
         run: |
@@ -358,7 +368,8 @@ jobs:
           UV_TEST_PUBLISH_CLOUDSMITH_TOKEN: ${{ secrets.UV_TEST_PUBLISH_CLOUDSMITH_TOKEN }}
           UV_TEST_PUBLISH_PYX_TOKEN: ${{ secrets.UV_TEST_PUBLISH_PYX_TOKEN }}
           UV_TEST_PUBLISH_PYTHON_VERSION: ${{ env.PYTHON_VERSION }}
-          UV_TEST_PUBLISH_GITLAB_OIDC_TOKEN: ${{ steps.load-gitlab-oidc-token.outputs.GITLAB_OIDC_TOKEN }}
+          UV_TEST_PUBLISH_GITLAB_PYPI_OIDC_TOKEN: ${{ steps.load-gitlab-oidc-token.outputs.GITLAB_PYPI_OIDC_TOKEN }}
+          UV_TEST_PUBLISH_GITLAB_PYX_OIDC_TOKEN: ${{ steps.load-gitlab-oidc-token.outputs.GITLAB_PYX_OIDC_TOKEN }}
 
   required-checks-passed:
     name: "all required jobs passed"

crates/uv-publish/src/lib.rs

@@ -40,7 +40,10 @@ use uv_redacted::{DisplaySafeUrl, DisplaySafeUrlError};
 use uv_warnings::warn_user;
 
 use crate::trusted_publishing::pypi::PyPIPublishingService;
-use crate::trusted_publishing::{TrustedPublishingError, TrustedPublishingToken};
+use crate::trusted_publishing::pyx::PyxPublishingService;
+use crate::trusted_publishing::{
+    TrustedPublishingError, TrustedPublishingService, TrustedPublishingToken,
+};
 
 #[derive(Error, Debug)]
 pub enum PublishError {
@@ -402,6 +405,7 @@ pub async fn check_trusted_publishing(
     username: Option<&str>,
     password: Option<&str>,
     keyring_provider: KeyringProviderType,
+    token_store: &PyxTokenStore,
     trusted_publishing: TrustedPublishing,
     registry: &DisplaySafeUrl,
     client: &BaseClient,
@@ -417,9 +421,21 @@ pub async fn check_trusted_publishing(
             }
 
             debug!("Attempting to get a token for trusted publishing");
+
             // Attempt to get a token for trusted publishing.
-            let service = PyPIPublishingService::new(registry, client);
-            match trusted_publishing::get_token(&service).await {
+            let token = if token_store.is_known_url(registry) {
+                debug!("Using trusted publishing flow for pyx");
+                PyxPublishingService::new(registry, client)
+                    .get_token()
+                    .await
+            } else {
+                debug!("Using trusted publishing flow for PyPI");
+                PyPIPublishingService::new(registry, client)
+                    .get_token()
+                    .await
+            };
+
+            match token {
                 // Success: we have a token for trusted publishing.
                 Ok(Some(token)) => Ok(TrustedPublishResult::Configured(token)),
                 // Failed to discover an ambient OIDC token.
@@ -447,11 +463,22 @@ pub async fn check_trusted_publishing(
                 return Err(PublishError::MixedCredentials(conflicts.join(" and ")));
             }
 
-            let service = PyPIPublishingService::new(registry, client);
-            let Some(token) = trusted_publishing::get_token(&service)
-                .await
-                .map_err(Box::new)?
-            else {
+            // Attempt to get a token for trusted publishing.
+            let token = if token_store.is_known_url(registry) {
+                debug!("Using trusted publishing flow for pyx");
+                PyxPublishingService::new(registry, client)
+                    .get_token()
+                    .await
+                    .map_err(Box::new)?
+            } else {
+                debug!("Using trusted publishing flow for PyPI");
+                PyPIPublishingService::new(registry, client)
+                    .get_token()
+                    .await
+                    .map_err(Box::new)?
+            };
+
+            let Some(token) = token else {
                 return Err(PublishError::TrustedPublishing(
                     TrustedPublishingError::NoToken.into(),
                 ));

crates/uv-publish/src/trusted_publishing.rs

@@ -12,6 +12,7 @@ use uv_redacted::{DisplaySafeUrl, DisplaySafeUrlError};
 use uv_static::EnvVars;
 
 pub(crate) mod pypi;
+pub(crate) mod pyx;
 
 #[derive(Debug, Error)]
 pub enum TrustedPublishingError {
@@ -35,12 +36,17 @@ pub enum TrustedPublishingError {
     #[error(transparent)]
     SerdeJson(#[from] serde_json::error::Error),
     #[error(
-        "PyPI returned error code {0}, is trusted publishing correctly configured?\nResponse: {1}\nToken claims, which must match the PyPI configuration: {2:#?}"
+        "Server returned error code {0}, is trusted publishing correctly configured?\nResponse: {1}\nToken claims, which must match the publisher configuration: {2:#?}"
     )]
-    Pypi(StatusCode, String, OidcTokenClaims),
+    TokenRejected(StatusCode, String, OidcTokenClaims),
     /// When trusted publishing is misconfigured, the error above should occur, not this one.
-    #[error("PyPI returned error code {0}, and the OIDC has an unexpected format.\nResponse: {1}")]
+    #[error(
+        "Server returned error code {0}, and the OIDC has an unexpected format.\nResponse: {1}"
+    )]
     InvalidOidcToken(StatusCode, String),
+    /// The user gave us a malformed upload URL for trusted publishing with pyx.
+    #[error("The upload URL `{0}` does not look like a valid pyx upload URL")]
+    InvalidPyxUploadUrl(DisplaySafeUrl),
 }
 
 #[derive(Deserialize)]
@@ -74,62 +80,88 @@ struct PublishToken {
 /// The payload of the OIDC token.
 #[derive(Deserialize, Debug)]
 #[allow(dead_code)]
-pub struct OidcTokenClaims {
+#[serde(untagged)]
+pub enum OidcTokenClaims {
+    GitHub(GitHubTokenClaims),
+    GitLab(GitLabTokenClaims),
+    Buildkite(BuildkiteTokenClaims),
+}
+
+/// The relevant payload of a GitHub OIDC token.
+#[derive(Deserialize, Debug)]
+#[allow(dead_code)]
+pub struct GitHubTokenClaims {
     sub: String,
     repository: String,
     repository_owner: String,
     repository_owner_id: String,
     job_workflow_ref: String,
     r#ref: String,
+    environment: Option<String>,
+}
+
+/// The relevant payload of a GitLab OIDC token.
+#[derive(Deserialize, Debug)]
+#[allow(dead_code)]
+pub struct GitLabTokenClaims {
+    sub: String,
+    project_path: String,
+    ci_config_ref_uri: String,
+    environment: Option<String>,
+}
+
+/// The relevant payload of a Buildkite OIDC token.
+#[derive(Deserialize, Debug)]
+#[allow(dead_code)]
+pub struct BuildkiteTokenClaims {
+    sub: String,
+    pipeline_slug: String,
+    organization_slug: String,
 }
 
 /// A service (i.e. uploadable index) that supports trusted publishing.
+///
+/// Interactions should go through the default [`get_token`]; implementors
+/// should implement the constituent trait methods.
 pub(crate) trait TrustedPublishingService {
-    /// Borrow the HTTP client with middleware.
+    /// Borrow an HTTP client with middleware.
     fn client(&self) -> &ClientWithMiddleware;
 
     /// Retrieve the service's expected OIDC audience.
     async fn audience(&self) -> Result<String, TrustedPublishingError>;
 
     /// Exchange an ambient OIDC identity token for a short-lived upload token on the service.
-    async fn publish_token(
+    async fn exchange_token(
         &self,
         oidc_token: ambient_id::IdToken,
     ) -> Result<TrustedPublishingToken, TrustedPublishingError>;
-}
 
-/// Returns the short-lived token to use for uploading.
-///
-/// Return states:
-/// - `Ok(Some(token))`: Successfully obtained a trusted publishing token.
-/// - `Ok(None)`: Not in a supported CI environment for trusted publishing.
-/// - `Err(...)`: An error occurred while trying to obtain the token.
-pub(crate) async fn get_token(
-    service: &impl TrustedPublishingService,
-) -> Result<Option<TrustedPublishingToken>, TrustedPublishingError> {
-    // Get the OIDC token's audience from the registry.
-    let audience = service.audience().await?;
-
-    // Perform ambient OIDC token discovery.
-    // Depending on the host (GitHub Actions, GitLab CI, etc.)
-    // this may perform additional network requests.
-    let oidc_token = get_oidc_token(&audience, service.client()).await?;
-
-    // Exchange the OIDC token for a short-lived upload token,
-    // if OIDC token discovery succeeded.
-    if let Some(oidc_token) = oidc_token {
-        let publish_token = service.publish_token(oidc_token).await?;
-
-        // If we're on GitHub Actions, mask the exchanged token in logs.
-        #[expect(clippy::print_stdout)]
-        if env::var(EnvVars::GITHUB_ACTIONS) == Ok("true".to_string()) {
-            println!("::add-mask::{publish_token}");
+    /// Perform the full trusted publishing token exchange.
+    async fn get_token(&self) -> Result<Option<TrustedPublishingToken>, TrustedPublishingError> {
+        // Get the OIDC token's audience from the registry.
+        let audience = self.audience().await?;
+
+        // Perform ambient OIDC token discovery.
+        // Depending on the host (GitHub Actions, GitLab CI, etc.)
+        // this may perform additional network requests.
+        let oidc_token = get_oidc_token(&audience, self.client()).await?;
+
+        // Exchange the OIDC token for a short-lived upload token,
+        // if OIDC token discovery succeeded.
+        if let Some(oidc_token) = oidc_token {
+            let publish_token = self.exchange_token(oidc_token).await?;
+
+            // If we're on GitHub Actions, mask the exchanged token in logs.
+            #[expect(clippy::print_stdout)]
+            if env::var(EnvVars::GITHUB_ACTIONS) == Ok("true".to_string()) {
+                println!("::add-mask::{publish_token}");
+            }
+
+            Ok(Some(publish_token))
+        } else {
+            // Not in a supported CI environment for trusted publishing.
+            Ok(None)
         }
-
-        Ok(Some(publish_token))
-    } else {
-        // Not in a supported CI environment for trusted publishing.
-        Ok(None)
     }
 }
 

crates/uv-publish/src/trusted_publishing/pypi.rs

@@ -61,7 +61,7 @@ impl TrustedPublishingService for PyPIPublishingService<'_> {
         Ok(audience.audience)
     }
 
-    async fn publish_token(
+    async fn exchange_token(
         &self,
         oidc_token: ambient_id::IdToken,
     ) -> Result<super::TrustedPublishingToken, super::TrustedPublishingError> {
@@ -107,7 +107,7 @@ impl TrustedPublishingService for PyPIPublishingService<'_> {
                     // configuration, so we're showing the body and the JWT claims for more context, see
                     // https://docs.pypi.org/trusted-publishers/troubleshooting/#token-minting
                     // for what the body can mean.
-                    Err(TrustedPublishingError::Pypi(
+                    Err(TrustedPublishingError::TokenRejected(
                         status,
                         String::from_utf8_lossy(&body).to_string(),
                         claims,