[auth] add workload identity token exchange [ci changed_files]

codex-rs/Cargo.lock

@@ -4537,6 +4537,22 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "codex-workload-identity"
+version = "0.0.0"
+dependencies = [
+ "codex-http-client",
+ "pretty_assertions",
+ "reqwest 0.12.28",
+ "serde",
+ "serde_json",
+ "tempfile",
+ "thiserror 2.0.18",
+ "tokio",
+ "url",
+ "wiremock",
+]
+
 [[package]]
 name = "color-eyre"
 version = "0.6.5"

codex-rs/Cargo.toml

@@ -91,6 +91,7 @@ members = [
     "tools",
     "v8-poc",
     "websocket-client",
+    "workload-identity",
     "utils/absolute-path",
     "utils/path-uri",
     "utils/cargo-bin",
@@ -169,6 +170,7 @@ codex-code-mode-protocol = { path = "code-mode-protocol" }
 codex-home = { path = "codex-home" }
 codex-http-client = { path = "http-client" }
 codex-websocket-client = { path = "websocket-client" }
+codex-workload-identity = { path = "workload-identity" }
 codex-config = { path = "config" }
 codex-connectors = { path = "connectors" }
 codex-connectors-extension = { path = "ext/connectors" }

codex-rs/workload-identity/BUILD.bazel

@@ -0,0 +1,6 @@
+load("//:defs.bzl", "codex_rust_crate")
+
+codex_rust_crate(
+    name = "workload-identity",
+    crate_name = "codex_workload_identity",
+)

codex-rs/workload-identity/Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+edition.workspace = true
+license.workspace = true
+name = "codex-workload-identity"
+version.workspace = true
+
+[lib]
+doctest = false
+name = "codex_workload_identity"
+path = "src/lib.rs"
+
+[lints]
+workspace = true
+
+[dependencies]
+codex-http-client = { workspace = true }
+reqwest = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
+thiserror = { workspace = true }
+tokio = { workspace = true, features = ["fs", "io-util", "sync", "time"] }
+url = { workspace = true }
+
+[dev-dependencies]
+pretty_assertions = { workspace = true }
+tempfile = { workspace = true }
+wiremock = { workspace = true }

codex-rs/workload-identity/src/assertion.rs

@@ -0,0 +1,50 @@
+use std::path::PathBuf;
+
+use tokio::io::AsyncReadExt;
+
+use crate::WorkloadIdentityError;
+
+const MAX_ASSERTION_BYTES: u64 = 16 * 1024;
+
+/// Source of the upstream assertion. File sources are reread for every exchange.
+#[derive(Clone)]
+pub enum WorkloadIdentityAssertionSource {
+    Environment(String),
+    File(PathBuf),
+}
+
+impl WorkloadIdentityAssertionSource {
+    pub(crate) async fn assertion(&self) -> Result<String, WorkloadIdentityError> {
+        let assertion = match self {
+            Self::Environment(assertion) => assertion.clone(),
+            Self::File(path) => {
+                let file = tokio::fs::File::open(path).await.map_err(|source| {
+                    WorkloadIdentityError::TokenFile {
+                        path: path.clone(),
+                        source,
+                    }
+                })?;
+                let mut bytes = Vec::new();
+                file.take(MAX_ASSERTION_BYTES + 1)
+                    .read_to_end(&mut bytes)
+                    .await
+                    .map_err(|source| WorkloadIdentityError::TokenFile {
+                        path: path.clone(),
+                        source,
+                    })?;
+                if bytes.len() as u64 > MAX_ASSERTION_BYTES {
+                    return Err(WorkloadIdentityError::AssertionTooLarge);
+                }
+                String::from_utf8(bytes).map_err(|_| WorkloadIdentityError::InvalidAssertion)?
+            }
+        };
+        let assertion = assertion.trim();
+        if assertion.len() as u64 > MAX_ASSERTION_BYTES {
+            return Err(WorkloadIdentityError::AssertionTooLarge);
+        }
+        if assertion.is_empty() || assertion.as_bytes().contains(&0) {
+            return Err(WorkloadIdentityError::InvalidAssertion);
+        }
+        Ok(assertion.to_string())
+    }
+}

codex-rs/workload-identity/src/exchange.rs

@@ -0,0 +1,327 @@
+use std::fmt;
+use std::time::Duration;
+use std::time::Instant;
+
+use codex_http_client::build_reqwest_client_with_custom_ca;
+use reqwest::Client;
+use reqwest::ClientBuilder;
+use reqwest::StatusCode;
+use reqwest::redirect::Policy;
+use serde::Deserialize;
+use tokio::sync::Mutex;
+use tokio::sync::Semaphore;
+use tokio::time::sleep;
+use url::Host;
+use url::Url;
+
+use crate::WorkloadIdentityConfig;
+use crate::WorkloadIdentityError;
+
+const DEFAULT_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
+pub(crate) const JWT_BEARER_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:jwt-bearer";
+const MAX_ACCESS_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
+const MAX_RESPONSE_BYTES: usize = 1024 * 1024;
+const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
+const RETRY_BASE_DELAY: Duration = Duration::from_millis(100);
+const RETRY_LIMIT: u32 = 2;
+const TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_WIF_TOKEN_URL_OVERRIDE";
+
+/// Exchanges assertions and retains only the current short-lived token in memory.
+pub struct WorkloadIdentityExchange {
+    cache: Mutex<CacheState>,
+    client: Client,
+    config: WorkloadIdentityConfig,
+    exchange_gate: Semaphore,
+    token_url: Url,
+}
+
+impl WorkloadIdentityExchange {
+    pub fn new(config: WorkloadIdentityConfig) -> Result<Self, WorkloadIdentityError> {
+        let (token_url, is_loopback_override) = token_url_from_environment()?;
+        let builder = Client::builder()
+            .redirect(Policy::none())
+            .timeout(REQUEST_TIMEOUT);
+        let builder = if is_loopback_override {
+            builder.no_proxy()
+        } else {
+            builder
+        };
+        Self::with_client_builder(config, token_url, builder)
+    }
+
+    pub(crate) fn with_client_builder(
+        config: WorkloadIdentityConfig,
+        token_url: Url,
+        builder: ClientBuilder,
+    ) -> Result<Self, WorkloadIdentityError> {
+        let client = build_reqwest_client_with_custom_ca(builder)
+            .map_err(|_| WorkloadIdentityError::HttpClientConfiguration)?;
+        Ok(Self {
+            cache: Mutex::new(CacheState::default()),
+            client,
+            config,
+            exchange_gate: Semaphore::new(1),
+            token_url,
+        })
+    }
+
+    /// Returns a cached token when possible and otherwise performs an exchange.
+    pub async fn resolve(&self) -> Result<WorkloadIdentityToken, WorkloadIdentityError> {
+        self.exchange(ExchangeMode::Resolve).await
+    }
+
+    /// Forces one fresh exchange after a downstream service rejects the current token.
+    pub async fn refresh(&self) -> Result<WorkloadIdentityToken, WorkloadIdentityError> {
+        self.exchange(ExchangeMode::Refresh).await
+    }
+
+    async fn exchange(
+        &self,
+        mode: ExchangeMode,
+    ) -> Result<WorkloadIdentityToken, WorkloadIdentityError> {
+        let now = Instant::now();
+        let observed_generation = {
+            let state = self.cache.lock().await;
+            if mode == ExchangeMode::Resolve
+                && let Some(cached) = &state.cached
+                && cached.refresh_at > now
+            {
+                return Ok(cached.token.clone());
+            }
+            state.generation
+        };
+        let _permit = self
+            .exchange_gate
+            .acquire()
+            .await
+            .map_err(|_| WorkloadIdentityError::ExchangeUnavailable)?;
+        let now = Instant::now();
+        let (fallback, next_generation) = {
+            let state = self.cache.lock().await;
+            if state.generation != observed_generation
+                && let Some(cached) = &state.cached
+            {
+                return Ok(cached.token.clone());
+            }
+            if mode == ExchangeMode::Resolve
+                && let Some(cached) = &state.cached
+                && cached.refresh_at > now
+            {
+                return Ok(cached.token.clone());
+            }
+            let fallback = match mode {
+                ExchangeMode::Refresh => None,
+                ExchangeMode::Resolve => state
+                    .cached
+                    .as_ref()
+                    .filter(|cached| cached.mandatory_refresh_at > now)
+                    .map(|cached| cached.token.clone()),
+            };
+            (fallback, state.generation.saturating_add(1))
+        };
+        let token = match self.exchange_uncached().await {
+            Ok(token) => token,
+            Err(error) => {
+                let Some(fallback) = fallback else {
+                    return Err(error);
+                };
+                let mut state = self.cache.lock().await;
+                if let Some(cached) = state.cached.as_mut() {
+                    cached.refresh_at =
+                        std::cmp::min(now + Duration::from_secs(30), cached.mandatory_refresh_at);
+                }
+                state.generation = next_generation;
+                return Ok(fallback);
+            }
+        };
+        let mut state = self.cache.lock().await;
+        state.generation = next_generation;
+        state.cached = Some(CachedToken::new(token.clone(), now));
+        Ok(token)
+    }
+
+    async fn exchange_uncached(&self) -> Result<WorkloadIdentityToken, WorkloadIdentityError> {
+        let assertion = self.config.assertion_source.assertion().await?;
+        let target = &self.config.target;
+        let form = [
+            ("grant_type", JWT_BEARER_GRANT_TYPE),
+            ("assertion", assertion.as_str()),
+            ("federation_rule_id", target.federation_rule_id.as_str()),
+            ("tenant_id", target.tenant_id.as_str()),
+            ("principal_id", target.principal_id.as_str()),
+            ("workspace_id", target.workspace_id.as_str()),
+        ];
+        let mut attempt = 0;
+        let mut response = loop {
+            match self
+                .client
+                .post(self.token_url.clone())
+                .form(&form)
+                .send()
+                .await
+            {
+                Ok(response) if is_retryable_status(response.status()) && attempt < RETRY_LIMIT => {
+                    attempt += 1;
+                    sleep(retry_delay(attempt)).await;
+                }
+                Ok(response) => break response,
+                Err(_) if attempt < RETRY_LIMIT => {
+                    attempt += 1;
+                    sleep(retry_delay(attempt)).await;
+                }
+                Err(_) => return Err(WorkloadIdentityError::ExchangeUnavailable),
+            }
+        };
+        if !response.status().is_success() {
+            return Err(WorkloadIdentityError::ExchangeRejected(
+                response.status().as_u16(),
+            ));
+        }
+        if response
+            .content_length()
+            .is_some_and(|length| length > MAX_RESPONSE_BYTES as u64)
+        {
+            return Err(WorkloadIdentityError::InvalidExchangeResponse);
+        }
+        let mut bytes = Vec::new();
+        while let Some(chunk) = response
+            .chunk()
+            .await
+            .map_err(|_| WorkloadIdentityError::InvalidExchangeResponse)?
+        {
+            if chunk.len() > MAX_RESPONSE_BYTES.saturating_sub(bytes.len()) {
+                return Err(WorkloadIdentityError::InvalidExchangeResponse);
+            }
+            bytes.extend_from_slice(&chunk);
+        }
+        let response: TokenExchangeResponse = serde_json::from_slice(&bytes)
+            .map_err(|_| WorkloadIdentityError::InvalidExchangeResponse)?;
+        response.into_token(&self.config.target)
+    }
+}
+
+fn retry_delay(attempt: u32) -> Duration {
+    RETRY_BASE_DELAY.saturating_mul(2_u32.saturating_pow(attempt.saturating_sub(1)))
+}
+
+fn is_retryable_status(status: StatusCode) -> bool {
+    status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error()
+}
+
+fn token_url_from_environment() -> Result<(Url, bool), WorkloadIdentityError> {
+    let Some(value) = std::env::var_os(TOKEN_URL_OVERRIDE_ENV_VAR) else {
+        return Url::parse(DEFAULT_TOKEN_URL)
+            .map(|url| (url, false))
+            .map_err(|_| WorkloadIdentityError::InvalidTokenUrl);
+    };
+    let value = value
+        .into_string()
+        .map_err(|_| WorkloadIdentityError::InvalidTokenUrl)?;
+    parse_loopback_token_url(&value).map(|url| (url, true))
+}
+
+pub(crate) fn parse_loopback_token_url(value: &str) -> Result<Url, WorkloadIdentityError> {
+    let url = Url::parse(value).map_err(|_| WorkloadIdentityError::InvalidTokenUrl)?;
+    if !url.username().is_empty()
+        || url.password().is_some()
+        || url.query().is_some()
+        || url.fragment().is_some()
+    {
+        return Err(WorkloadIdentityError::InvalidTokenUrl);
+    }
+    let loopback = match url.host() {
+        Some(Host::Domain(domain)) => domain.eq_ignore_ascii_case("localhost"),
+        Some(Host::Ipv4(address)) => address.is_loopback(),
+        Some(Host::Ipv6(address)) => address.is_loopback(),
+        None => false,
+    };
+    if !matches!(url.scheme(), "http" | "https") || !loopback {
+        return Err(WorkloadIdentityError::InvalidTokenUrl);
+    }
+    Ok(url)
+}
+
+#[derive(Default)]
+struct CacheState {
+    cached: Option<CachedToken>,
+    generation: u64,
+}
+
+struct CachedToken {
+    mandatory_refresh_at: Instant,
+    refresh_at: Instant,
+    token: WorkloadIdentityToken,
+}
+
+impl CachedToken {
+    fn new(token: WorkloadIdentityToken, now: Instant) -> Self {
+        let lifetime = Duration::from_secs(token.expires_in);
+        let advisory_margin = std::cmp::min(Duration::from_secs(120), lifetime / 2);
+        let mandatory_margin = std::cmp::min(Duration::from_secs(30), lifetime / 4);
+        Self {
+            mandatory_refresh_at: now + lifetime.saturating_sub(mandatory_margin),
+            refresh_at: now + lifetime.saturating_sub(advisory_margin),
+            token,
+        }
+    }
+}
+
+#[derive(Clone, PartialEq, Eq)]
+pub struct WorkloadIdentityToken {
+    pub access_token: String,
+    pub chatgpt_account_id: String,
+    pub chatgpt_plan_type: Option<String>,
+    pub expires_in: u64,
+}
+
+impl fmt::Debug for WorkloadIdentityToken {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter
+            .debug_struct("WorkloadIdentityToken")
+            .field("access_token", &"[redacted]")
+            .field("chatgpt_account_id", &self.chatgpt_account_id)
+            .field("chatgpt_plan_type", &self.chatgpt_plan_type)
+            .field("expires_in", &self.expires_in)
+            .finish()
+    }
+}
+
+#[derive(Deserialize)]
+struct TokenExchangeResponse {
+    access_token: String,
+    chatgpt_account_id: String,
+    chatgpt_plan_type: Option<String>,
+    expires_in: u64,
+    token_type: String,
+    user_id: String,
+}
+
+impl TokenExchangeResponse {
+    fn into_token(
+        self,
+        expected: &crate::WorkloadIdentityTarget,
+    ) -> Result<WorkloadIdentityToken, WorkloadIdentityError> {
+        let lifetime = Duration::from_secs(self.expires_in);
+        if self.access_token.trim().is_empty()
+            || !self.token_type.eq_ignore_ascii_case("bearer")
+            || lifetime.is_zero()
+            || lifetime > MAX_ACCESS_TOKEN_LIFETIME
+            || self.chatgpt_account_id != expected.workspace_id
+            || self.user_id != expected.principal_id
+        {
+            return Err(WorkloadIdentityError::InvalidExchangeResponse);
+        }
+        Ok(WorkloadIdentityToken {
+            access_token: self.access_token,
+            chatgpt_account_id: self.chatgpt_account_id,
+            chatgpt_plan_type: self.chatgpt_plan_type,
+            expires_in: self.expires_in,
+        })
+    }
+}
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum ExchangeMode {
+    Resolve,
+    Refresh,
+}

codex-rs/workload-identity/src/lib.rs

@@ -0,0 +1,88 @@
+mod assertion;
+mod exchange;
+
+use std::path::PathBuf;
+
+pub use assertion::WorkloadIdentityAssertionSource;
+pub use exchange::WorkloadIdentityExchange;
+pub use exchange::WorkloadIdentityToken;
+use thiserror::Error;
+
+pub const FEDERATION_RULE_ID_ENV_VAR: &str = "OPENAI_FEDERATION_RULE_ID";
+pub const IDENTITY_TOKEN_ENV_VAR: &str = "OPENAI_IDENTITY_TOKEN";
+pub const IDENTITY_TOKEN_FILE_ENV_VAR: &str = "OPENAI_IDENTITY_TOKEN_FILE";
+pub const PRINCIPAL_ID_ENV_VAR: &str = "OPENAI_PRINCIPAL_ID";
+pub const TENANT_ID_ENV_VAR: &str = "OPENAI_TENANT_ID";
+pub const WORKSPACE_ID_ENV_VAR: &str = "OPENAI_WORKSPACE_ID";
+
+/// Identifies the pre-provisioned OpenAI principal selected by a federation rule.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct WorkloadIdentityTarget {
+    pub federation_rule_id: String,
+    pub principal_id: String,
+    pub tenant_id: String,
+    pub workspace_id: String,
+}
+
+/// Complete input for exchanging an upstream assertion for ChatGPT auth.
+#[derive(Clone)]
+pub struct WorkloadIdentityConfig {
+    pub(crate) assertion_source: WorkloadIdentityAssertionSource,
+    pub(crate) target: WorkloadIdentityTarget,
+}
+
+impl WorkloadIdentityConfig {
+    pub fn new(
+        target: WorkloadIdentityTarget,
+        assertion_source: WorkloadIdentityAssertionSource,
+    ) -> Result<Self, WorkloadIdentityError> {
+        let target = WorkloadIdentityTarget {
+            federation_rule_id: normalized_field(target.federation_rule_id, "federation_rule_id")?,
+            principal_id: normalized_field(target.principal_id, "principal_id")?,
+            tenant_id: normalized_field(target.tenant_id, "tenant_id")?,
+            workspace_id: normalized_field(target.workspace_id, "workspace_id")?,
+        };
+        Ok(Self {
+            assertion_source,
+            target,
+        })
+    }
+}
+
+fn normalized_field(value: String, name: &'static str) -> Result<String, WorkloadIdentityError> {
+    let value = value.trim();
+    if value.is_empty() {
+        return Err(WorkloadIdentityError::InvalidConfigurationField(name));
+    }
+    Ok(value.to_string())
+}
+
+#[derive(Debug, Error)]
+pub enum WorkloadIdentityError {
+    #[error("workload identity field {0} must not be empty")]
+    InvalidConfigurationField(&'static str),
+    #[error("the workload identity assertion is invalid")]
+    InvalidAssertion,
+    #[error("the workload identity assertion exceeds 16 KiB")]
+    AssertionTooLarge,
+    #[error("could not read workload identity token file {path}")]
+    TokenFile {
+        path: PathBuf,
+        #[source]
+        source: std::io::Error,
+    },
+    #[error("could not configure the workload identity HTTP client")]
+    HttpClientConfiguration,
+    #[error("CODEX_WIF_TOKEN_URL_OVERRIDE must use loopback HTTP(S)")]
+    InvalidTokenUrl,
+    #[error("the workload identity token exchange is unavailable")]
+    ExchangeUnavailable,
+    #[error("the workload identity token exchange was rejected with HTTP {0}")]
+    ExchangeRejected(u16),
+    #[error("the workload identity token exchange returned an invalid response")]
+    InvalidExchangeResponse,
+}
+
+#[cfg(test)]
+#[path = "tests.rs"]
+mod tests;

codex-rs/workload-identity/src/tests.rs

@@ -0,0 +1,277 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
+use pretty_assertions::assert_eq;
+use reqwest::Client;
+use tempfile::TempDir;
+use url::Url;
+use wiremock::Mock;
+use wiremock::MockServer;
+use wiremock::ResponseTemplate;
+use wiremock::matchers::method;
+use wiremock::matchers::path;
+
+use super::*;
+use crate::exchange::JWT_BEARER_GRANT_TYPE;
+use crate::exchange::parse_loopback_token_url;
+
+fn target() -> WorkloadIdentityTarget {
+    WorkloadIdentityTarget {
+        federation_rule_id: "rule-one".to_string(),
+        principal_id: "user-one".to_string(),
+        tenant_id: "tenant-one".to_string(),
+        workspace_id: "workspace-one".to_string(),
+    }
+}
+
+fn config(source: WorkloadIdentityAssertionSource) -> WorkloadIdentityConfig {
+    WorkloadIdentityConfig::new(target(), source).expect("valid workload identity config")
+}
+
+fn make_exchange(
+    server: &MockServer,
+    source: WorkloadIdentityAssertionSource,
+) -> WorkloadIdentityExchange {
+    WorkloadIdentityExchange::with_client_builder(
+        config(source),
+        Url::parse(&format!("{}/oauth/token", server.uri())).expect("valid token URL"),
+        Client::builder().no_proxy(),
+    )
+    .expect("valid exchange")
+}
+
+fn success(access_token: &str) -> ResponseTemplate {
+    ResponseTemplate::new(200).set_body_json(serde_json::json!({
+        "access_token": access_token,
+        "token_type": "Bearer",
+        "expires_in": 600,
+        "chatgpt_account_id": "workspace-one",
+        "chatgpt_plan_type": "enterprise",
+        "user_id": "user-one"
+    }))
+}
+
+#[test]
+fn configuration_rejects_empty_selectors() {
+    let mut target = target();
+    target.workspace_id = "  ".to_string();
+    assert!(matches!(
+        WorkloadIdentityConfig::new(
+            target,
+            WorkloadIdentityAssertionSource::Environment("assertion".to_string())
+        ),
+        Err(WorkloadIdentityError::InvalidConfigurationField(
+            "workspace_id"
+        ))
+    ));
+}
+
+#[test]
+fn loopback_override_rejects_non_loopback_and_url_metadata() {
+    for valid in [
+        "https://localhost:3000/oauth/token",
+        "http://127.0.0.1:3000/oauth/token",
+        "http://[::1]:3000/oauth/token",
+    ] {
+        assert!(parse_loopback_token_url(valid).is_ok());
+    }
+    for invalid in [
+        "https://auth.example.com/oauth/token",
+        "http://auth.localhost:3000/oauth/token",
+        "https://user:password@localhost/oauth/token",
+        "https://localhost/oauth/token?assertion=secret",
+        "file:///tmp/token",
+    ] {
+        assert!(matches!(
+            parse_loopback_token_url(invalid),
+            Err(WorkloadIdentityError::InvalidTokenUrl)
+        ));
+    }
+}
+
+#[tokio::test]
+async fn exchange_sends_the_rfc_7523_contract_and_caches_the_result() {
+    let server = MockServer::start().await;
+    Mock::given(method("POST"))
+        .and(path("/oauth/token"))
+        .respond_with(success("access-one"))
+        .mount(&server)
+        .await;
+    let exchange = make_exchange(
+        &server,
+        WorkloadIdentityAssertionSource::Environment("assertion-one".to_string()),
+    );
+
+    let expected = WorkloadIdentityToken {
+        access_token: "access-one".to_string(),
+        chatgpt_account_id: "workspace-one".to_string(),
+        chatgpt_plan_type: Some("enterprise".to_string()),
+        expires_in: 600,
+    };
+    assert_eq!(exchange.resolve().await.expect("first exchange"), expected);
+    assert_eq!(exchange.resolve().await.expect("cached exchange"), expected);
+
+    let requests = server.received_requests().await.expect("received requests");
+    assert_eq!(requests.len(), 1);
+    assert_eq!(
+        url::form_urlencoded::parse(&requests[0].body)
+            .into_owned()
+            .collect::<HashMap<_, _>>(),
+        HashMap::from([
+            ("assertion".to_string(), "assertion-one".to_string()),
+            ("federation_rule_id".to_string(), "rule-one".to_string()),
+            ("grant_type".to_string(), JWT_BEARER_GRANT_TYPE.to_string()),
+            ("principal_id".to_string(), "user-one".to_string()),
+            ("tenant_id".to_string(), "tenant-one".to_string()),
+            ("workspace_id".to_string(), "workspace-one".to_string()),
+        ])
+    );
+}
+
+#[tokio::test]
+async fn file_source_is_reread_on_refresh() {
+    let server = MockServer::start().await;
+    Mock::given(method("POST"))
+        .and(path("/oauth/token"))
+        .respond_with(success("access-one"))
+        .mount(&server)
+        .await;
+    let temp_dir = TempDir::new().expect("tempdir");
+    let token_file = temp_dir.path().join("identity-token");
+    tokio::fs::write(&token_file, "assertion-one\n")
+        .await
+        .expect("write assertion");
+    let exchange = make_exchange(
+        &server,
+        WorkloadIdentityAssertionSource::File(token_file.clone()),
+    );
+
+    exchange.resolve().await.expect("initial exchange");
+    tokio::fs::write(&token_file, "assertion-two\n")
+        .await
+        .expect("rotate assertion");
+    exchange.refresh().await.expect("refresh exchange");
+
+    let requests = server.received_requests().await.expect("received requests");
+    let assertions = requests
+        .iter()
+        .map(|request| {
+            url::form_urlencoded::parse(&request.body)
+                .find(|(name, _)| name == "assertion")
+                .map(|(_, value)| value.into_owned())
+                .expect("assertion field")
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(assertions, vec!["assertion-one", "assertion-two"]);
+}
+
+#[tokio::test]
+async fn concurrent_resolution_performs_one_exchange() {
+    let server = MockServer::start().await;
+    Mock::given(method("POST"))
+        .respond_with(success("access-one").set_delay(Duration::from_millis(50)))
+        .mount(&server)
+        .await;
+    let exchange = Arc::new(make_exchange(
+        &server,
+        WorkloadIdentityAssertionSource::Environment("assertion-one".to_string()),
+    ));
+
+    let tasks = (0..8)
+        .map(|_| {
+            let exchange = Arc::clone(&exchange);
+            tokio::spawn(async move { exchange.resolve().await })
+        })
+        .collect::<Vec<_>>();
+    for task in tasks {
+        task.await.expect("join exchange").expect("exchange");
+    }
+    assert_eq!(server.received_requests().await.expect("requests").len(), 1);
+
+    let refreshes = (0..8)
+        .map(|_| {
+            let exchange = Arc::clone(&exchange);
+            tokio::spawn(async move { exchange.refresh().await })
+        })
+        .collect::<Vec<_>>();
+    for refresh in refreshes {
+        refresh.await.expect("join refresh").expect("refresh");
+    }
+    assert_eq!(server.received_requests().await.expect("requests").len(), 2);
+}
+
+#[tokio::test]
+async fn exchange_retries_transient_statuses_without_exposing_secrets() {
+    let server = MockServer::start().await;
+    let calls = Arc::new(AtomicUsize::new(0));
+    Mock::given(method("POST"))
+        .respond_with({
+            let calls = Arc::clone(&calls);
+            move |_request: &wiremock::Request| match calls.fetch_add(1, Ordering::SeqCst) {
+                0 => ResponseTemplate::new(503).set_body_string("sensitive detail"),
+                1 => ResponseTemplate::new(429),
+                _ => success("access-one"),
+            }
+        })
+        .mount(&server)
+        .await;
+    let exchange = make_exchange(
+        &server,
+        WorkloadIdentityAssertionSource::Environment("sensitive-assertion".to_string()),
+    );
+
+    assert_eq!(
+        exchange
+            .resolve()
+            .await
+            .expect("retried exchange")
+            .access_token,
+        "access-one"
+    );
+    assert_eq!(calls.load(Ordering::SeqCst), 3);
+
+    let rejected_server = MockServer::start().await;
+    Mock::given(method("POST"))
+        .respond_with(ResponseTemplate::new(400).set_body_string("sensitive detail"))
+        .mount(&rejected_server)
+        .await;
+    let rejected = make_exchange(
+        &rejected_server,
+        WorkloadIdentityAssertionSource::Environment("sensitive-assertion".to_string()),
+    )
+    .resolve()
+    .await
+    .expect_err("exchange should be rejected")
+    .to_string();
+    assert_eq!(
+        rejected,
+        "the workload identity token exchange was rejected with HTTP 400"
+    );
+}
+
+#[tokio::test]
+async fn exchange_rejects_mismatched_or_overlong_tokens() {
+    let server = MockServer::start().await;
+    Mock::given(method("POST"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
+            "access_token": "access-one",
+            "token_type": "Bearer",
+            "expires_in": 3601,
+            "chatgpt_account_id": "workspace-two",
+            "user_id": "user-two"
+        })))
+        .mount(&server)
+        .await;
+    let exchange = make_exchange(
+        &server,
+        WorkloadIdentityAssertionSource::Environment("assertion-one".to_string()),
+    );
+
+    assert!(matches!(
+        exchange.resolve().await,
+        Err(WorkloadIdentityError::InvalidExchangeResponse)
+    ));
+}