yapi_ai/auth/
federation.rs

1//! Anthropic workload identity federation: exchanges an external OpenID
2//! Connect (OIDC) token for a short-lived Anthropic access token and caches
3//! it per configuration.
4//!
5//! Port of the Anthropic SDK's `oidcFederationProvider` and `TokenCache`
6//! (`@anthropic-ai/sdk` 0.124.0), which pi `v1.0.0` uses when the
7//! `ANTHROPIC_FEDERATION_RULE_ID`, `ANTHROPIC_ORGANIZATION_ID` and
8//! `ANTHROPIC_IDENTITY_TOKEN_FILE` provider settings are set.
9
10use std::collections::HashMap;
11use std::sync::{Mutex, OnceLock};
12
13use serde_json::{Map, Value};
14use tokio_util::sync::CancellationToken;
15
16use crate::credentials::{ProviderEnv, provider_env_value};
17
18const GRANT_TYPE_JWT_BEARER: &str = "urn:ietf:params:oauth:grant-type:jwt-bearer";
19const TOKEN_ENDPOINT: &str = "/v1/oauth/token";
20/// The `anthropic-beta` value that requests with a federated bearer token send.
21pub const OAUTH_API_BETA: &str = "oauth-2025-04-20";
22const FEDERATION_BETA: &str = "oidc-federation-2026-04-01";
23const SDK_VERSION: &str = "0.124.0";
24/// A cached token with less validity left is exchanged again.
25const MANDATORY_REFRESH_SECONDS: u64 = 30;
26const MAX_ASSERTION_BYTES: usize = 16 * 1024;
27const MAX_ERROR_BODY_CHARS: usize = 2000;
28
29/// What the exchange needs, from the provider settings.
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub struct Config {
32    base_url: String,
33    federation_rule_id: String,
34    organization_id: String,
35    identity_token_file: String,
36    service_account_id: Option<String>,
37    workspace_id: Option<String>,
38}
39
40impl Config {
41    /// The federation configuration for a request to `base_url`, when the
42    /// three required settings are present.
43    pub fn from_env(base_url: &str, env: Option<&ProviderEnv>) -> Option<Config> {
44        let value = |name: &str| provider_env_value(name, env);
45        Some(Config {
46            base_url: base_url.trim_end_matches('/').to_owned(),
47            federation_rule_id: value("ANTHROPIC_FEDERATION_RULE_ID")?,
48            organization_id: value("ANTHROPIC_ORGANIZATION_ID")?,
49            identity_token_file: value("ANTHROPIC_IDENTITY_TOKEN_FILE")?,
50            service_account_id: value("ANTHROPIC_SERVICE_ACCOUNT_ID"),
51            workspace_id: value("ANTHROPIC_WORKSPACE_ID"),
52        })
53    }
54
55    fn key(&self) -> String {
56        format!("{self:?}")
57    }
58}
59
60struct Cached {
61    token: String,
62    expires_at: u64,
63}
64
65fn cache() -> &'static Mutex<HashMap<String, Cached>> {
66    static CACHE: OnceLock<Mutex<HashMap<String, Cached>>> = OnceLock::new();
67    CACHE.get_or_init(Default::default)
68}
69
70fn now_seconds() -> u64 {
71    super::now_ms() / 1000
72}
73
74/// An access token for `config`: the cached one while it has more than 30s
75/// left, else a fresh exchange. `force` skips the cache, as after a
76/// 401.
77pub async fn token(
78    config: &Config,
79    force: bool,
80    cancel: &CancellationToken,
81) -> Result<String, String> {
82    if !force
83        && let Ok(cache) = cache().lock()
84        && let Some(cached) = cache.get(&config.key())
85        && cached.expires_at > now_seconds() + MANDATORY_REFRESH_SECONDS
86    {
87        return Ok(cached.token.clone());
88    }
89    let (token, expires_in) = exchange(config, cancel).await?;
90    if let Ok(mut cache) = cache().lock() {
91        cache.insert(
92            config.key(),
93            Cached {
94                token: token.clone(),
95                expires_at: now_seconds() + expires_in,
96            },
97        );
98    }
99    Ok(token)
100}
101
102/// Drops the cached token for `config`.
103pub fn invalidate(config: &Config) {
104    if let Ok(mut cache) = cache().lock() {
105        cache.remove(&config.key());
106    }
107}
108
109/// The SDK's `requireSecureTokenEndpoint`: the assertion is sent only over
110/// HTTPS, or plain HTTP to a loopback host.
111fn require_secure(base_url: &str) -> Result<(), String> {
112    let url = url::Url::parse(base_url)
113        .map_err(|err| format!("Invalid token endpoint base URL \"{base_url}\": {err}"))?;
114    if url.scheme() == "https" {
115        return Ok(());
116    }
117    let host = url
118        .host_str()
119        .unwrap_or_default()
120        .trim_start_matches('[')
121        .trim_end_matches(']')
122        .to_lowercase();
123    if url.scheme() == "http" && matches!(host.as_str(), "localhost" | "127.0.0.1" | "::1") {
124        return Ok(());
125    }
126    Err(format!(
127        "Refusing to send credential over non-https token endpoint \"{base_url}\""
128    ))
129}
130
131/// The SDK's `redactSensitive`: JSON error bodies keep only RFC 6749 error
132/// fields. Other text is truncated.
133fn redact(body: &str) -> String {
134    match serde_json::from_str::<Value>(body) {
135        Ok(Value::Object(fields)) => {
136            let safe: Map<String, Value> = fields
137                .into_iter()
138                .filter(|(key, _)| {
139                    matches!(key.as_str(), "error" | "error_description" | "error_uri")
140                })
141                .collect();
142            yapi_types::json::stringify(&Value::Object(safe))
143        }
144        Ok(_) => "null".to_owned(),
145        Err(_) if body.chars().count() <= MAX_ERROR_BODY_CHARS => body.to_owned(),
146        Err(_) => {
147            let total = body.chars().count();
148            let head: String = body.chars().take(MAX_ERROR_BODY_CHARS).collect();
149            format!("{head}... <{} more chars>", total - MAX_ERROR_BODY_CHARS)
150        }
151    }
152}
153
154async fn exchange(config: &Config, cancel: &CancellationToken) -> Result<(String, u64), String> {
155    require_secure(&config.base_url)?;
156    let path = &config.identity_token_file;
157    let content = tokio::fs::read_to_string(path)
158        .await
159        .map_err(|err| format!("Failed to read identity token file at {path}: {err}"))?;
160    let assertion = content.trim();
161    if assertion.is_empty() {
162        return Err(format!("Identity token file at {path} is empty"));
163    }
164    if assertion.len() > MAX_ASSERTION_BYTES {
165        return Err(format!(
166            "Identity token is {} KiB, exceeds the 16 KiB assertion limit",
167            assertion.len().div_ceil(1024)
168        ));
169    }
170    let mut body = Map::new();
171    body.insert("grant_type".into(), GRANT_TYPE_JWT_BEARER.into());
172    body.insert("assertion".into(), assertion.into());
173    body.insert(
174        "federation_rule_id".into(),
175        config.federation_rule_id.clone().into(),
176    );
177    body.insert(
178        "organization_id".into(),
179        config.organization_id.clone().into(),
180    );
181    if let Some(id) = &config.service_account_id {
182        body.insert("service_account_id".into(), id.clone().into());
183    }
184    if let Some(id) = &config.workspace_id {
185        body.insert("workspace_id".into(), id.clone().into());
186    }
187    let url = format!("{}{TOKEN_ENDPOINT}", config.base_url);
188    let request = crate::http::client()
189        .post(&url)
190        .header("Content-Type", "application/json")
191        .header(
192            "anthropic-beta",
193            format!("{OAUTH_API_BETA},{FEDERATION_BETA}"),
194        )
195        .header(
196            "User-Agent",
197            format!("anthropic-sdk-typescript/{SDK_VERSION} oidcFederationProvider"),
198        )
199        .body(yapi_types::json::stringify(&Value::Object(body)));
200    let response = super::send(request, cancel)
201        .await
202        .map_err(|err| match err {
203            super::AuthError::Cancelled => crate::http::ABORTED_BEFORE_RESPONSE.to_owned(),
204            super::AuthError::Failed(_) => {
205                format!("Failed to reach token endpoint {url}: TypeError: fetch failed")
206            }
207        })?;
208    let status = response.status().as_u16();
209    let request_id = response
210        .headers()
211        .get("Request-Id")
212        .and_then(|value| value.to_str().ok())
213        .map(str::to_owned);
214    let text = response.text().await.unwrap_or_default();
215    if !(200..300).contains(&status) {
216        let mut hint = String::new();
217        if status == 401 {
218            let workspace = if config.workspace_id.is_some() {
219                ""
220            } else {
221                "If your federation rule is scoped to multiple workspaces, set the ANTHROPIC_WORKSPACE_ID environment variable, the 'workspace_id' config key, or the `workspaceId` option. "
222            };
223            hint = format!(
224                " Ensure your federation rule matches your identity token. {workspace}View your authentication events in the Workload identity page of Claude Console for more details."
225            );
226        }
227        let id = request_id
228            .map(|id| format!(" (request-id {id})"))
229            .unwrap_or_default();
230        return Err(format!(
231            "Token exchange failed with status {status}{id}: {}{hint}",
232            redact(&text)
233        ));
234    }
235    let data: Value = serde_json::from_str(&text)
236        .map_err(|_| format!("Token endpoint returned non-JSON response (status {status})"))?;
237    let token = data["access_token"]
238        .as_str()
239        .filter(|token| !token.is_empty())
240        .ok_or_else(|| {
241            format!(
242                "Token endpoint response missing access_token: {}",
243                redact(&text)
244            )
245        })?;
246    if let Some(kind) = data["token_type"].as_str()
247        && !kind.is_empty()
248        && !kind.eq_ignore_ascii_case("bearer")
249    {
250        return Err(format!(
251            "Token endpoint response: unsupported token_type \"{kind}\" (want Bearer)"
252        ));
253    }
254    let expires_in = match &data["expires_in"] {
255        Value::Number(number) => number.as_f64(),
256        Value::String(text) => text.trim().parse::<f64>().ok(),
257        _ => None,
258    }
259    .filter(|seconds| seconds.is_finite())
260    .ok_or_else(|| {
261        format!(
262            "Token endpoint response missing required fields: {}",
263            redact(&text)
264        )
265    })?;
266    Ok((token.to_owned(), expires_in.max(0.0) as u64))
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    #[test]
274    fn reads_settings_guards_endpoint() {
275        let mut env = ProviderEnv::new();
276        env.insert("ANTHROPIC_FEDERATION_RULE_ID".into(), "rule".into());
277        env.insert("ANTHROPIC_ORGANIZATION_ID".into(), "org".into());
278        assert_eq!(
279            Config::from_env("https://api.anthropic.com", Some(&env)),
280            None
281        );
282        env.insert("ANTHROPIC_IDENTITY_TOKEN_FILE".into(), "/t".into());
283        let config = Config::from_env("https://api.anthropic.com/", Some(&env)).unwrap();
284        assert_eq!(config.base_url, "https://api.anthropic.com");
285        assert!(require_secure("http://127.0.0.1:9").is_ok());
286        assert!(require_secure("http://[::1]:9").is_ok());
287        assert_eq!(
288            require_secure("http://example.com"),
289            Err(
290                "Refusing to send credential over non-https token endpoint \"http://example.com\""
291                    .into()
292            )
293        );
294        assert_eq!(
295            redact(r#"{"error":"invalid_grant","assertion":"secret"}"#),
296            r#"{"error":"invalid_grant"}"#
297        );
298    }
299}
yapi_ai/auth/federation.rs

yapi_ai/auth/
federation.rs
