diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 1986b0b23f..a0a07e91e0 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -5dfdf811c862c689e42f03f485e02f9f1e7bc7ca +747336910e5c7498544613ad7863cea59f4db6a1 diff --git a/backend/windmill-ai/src/ai_providers.rs b/backend/windmill-ai/src/ai_providers.rs index 311c8d3eb7..489a965ef1 100644 --- a/backend/windmill-ai/src/ai_providers.rs +++ b/backend/windmill-ai/src/ai_providers.rs @@ -213,7 +213,9 @@ impl AIProvider { AIProvider::Groq => Ok("https://api.groq.com/openai/v1".to_string()), AIProvider::OpenRouter => Ok("https://openrouter.ai/api/v1".to_string()), AIProvider::TogetherAI => Ok("https://api.together.xyz/v1".to_string()), - AIProvider::Anthropic => Ok("https://api.anthropic.com/v1".to_string()), + AIProvider::Anthropic => { + Ok(crate::providers::anthropic::ANTHROPIC_API_BASE_URL.to_string()) + } AIProvider::Mistral => Ok("https://api.mistral.ai/v1".to_string()), AIProvider::TypeSafe => Ok(TYPESAFE_BASE_URL.to_string()), AIProvider::Cloudflare => Err(Error::BadRequest( diff --git a/backend/windmill-ai/src/credentials.rs b/backend/windmill-ai/src/credentials.rs index 5a779938ed..f5489c5c30 100644 --- a/backend/windmill-ai/src/credentials.rs +++ b/backend/windmill-ai/src/credentials.rs @@ -32,6 +32,21 @@ pub struct ProviderCredentials { /// A role trust policy that conditions on `sub` has to admit both, or only /// one of the copilot and the AI agent step will be able to assume the role. pub oidc_role_arn: Option, + /// Anthropic Workload Identity Federation settings, from an `anthropic_federated` + /// resource. The API proxy exchanges an OIDC token minted for the requesting + /// user for a short-lived access token, which then goes out as `access_token`. + pub anthropic_federation: Option, pub platform: AIPlatform, pub custom_headers: HashMap, } + +/// The identifiers a Claude Console federation rule needs in a token exchange. +/// None of them is a secret: what authorises the exchange is the OIDC token. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct AnthropicFederation { + pub federation_rule_id: String, + pub organization_id: String, + pub service_account_id: String, + /// Required by Anthropic only when the rule is enabled in several workspaces. + pub workspace_id: Option, +} diff --git a/backend/windmill-ai/src/providers/anthropic.rs b/backend/windmill-ai/src/providers/anthropic.rs index 7acd2db939..2a3417bc62 100644 --- a/backend/windmill-ai/src/providers/anthropic.rs +++ b/backend/windmill-ai/src/providers/anthropic.rs @@ -2,6 +2,7 @@ use super::{anthropic_model_rejects_sampling_params, REASONING_OFF_SENTINEL}; use crate::{ ai_google::parse_data_url, ai_providers::{AIPlatform, AIProvider, DISABLE_ANTHROPIC_PROMPT_CACHING}, + credentials::AnthropicFederation, image_handler::prepare_messages_for_api, proxy::{ add_user_to_body, common_outbound_headers, credential_header, ProxyBuildArgs, ProxyRequest, @@ -837,6 +838,169 @@ impl QueryBuilder for AnthropicQueryBuilder { } } +pub const ANTHROPIC_API_BASE_URL: &str = "https://api.anthropic.com/v1"; + +/// Audience of the OIDC token Windmill presents in an Anthropic token exchange, +/// which a federation rule can pin with its `audience` matcher. +pub const ANTHROPIC_OIDC_AUDIENCE: &str = "https://api.anthropic.com"; + +/// Lifetime of that OIDC token. A federation issuer refuses a token that lives +/// longer than its maximum (one hour by default), and Anthropic caps the access +/// token at twice the assertion's remaining life, so 30 minutes still lets a rule +/// mint its default one-hour token. +pub const ANTHROPIC_OIDC_TOKEN_LIFETIME_SECS: i64 = 1800; + +/// Exchange again this long before the access token expires, so a long streamed +/// response that starts on a reused token cannot outlive it. Capped at half the +/// token's lifetime: a rule may mint tokens for as little as a minute, and a +/// margin longer than that would exchange again on every request. +const FEDERATED_TOKEN_REFRESH_MARGIN: std::time::Duration = std::time::Duration::from_secs(300); + +/// A short-lived Anthropic access token minted by a federated token exchange. +#[derive(Clone, Debug)] +pub struct FederatedAccessToken { + pub access_token: String, + refresh_at: std::time::SystemTime, +} + +impl FederatedAccessToken { + fn new(access_token: String, lifetime: std::time::Duration) -> Self { + // Anthropic mints tokens for at most a day; the clamp keeps an out-of-range + // `expires_in` from overflowing the clock arithmetic. + let lifetime = lifetime.min(std::time::Duration::from_secs(86400)); + let margin = FEDERATED_TOKEN_REFRESH_MARGIN.min(lifetime / 2); + Self { access_token, refresh_at: std::time::SystemTime::now() + lifetime - margin } + } + + pub fn is_fresh(&self) -> bool { + std::time::SystemTime::now() < self.refresh_at + } +} + +#[derive(Serialize)] +struct FederatedTokenRequest<'a> { + grant_type: &'static str, + assertion: &'a str, + federation_rule_id: &'a str, + organization_id: &'a str, + service_account_id: &'a str, + #[serde(skip_serializing_if = "Option::is_none")] + workspace_id: Option<&'a str>, +} + +#[derive(Deserialize)] +struct FederatedTokenResponse { + access_token: String, + expires_in: u64, +} + +/// Exchange an OIDC token for an Anthropic access token at `POST /v1/oauth/token` +/// (RFC 7523 `jwt-bearer` grant). `client` must already be pinned to the +/// validated address of `base_url`. +pub async fn exchange_federated_token( + client: &reqwest::Client, + base_url: &str, + federation: &AnthropicFederation, + assertion: &str, +) -> Result { + let exchanged = request_federated_token(client, base_url, federation, assertion).await; + windmill_common::feature_usage::log_feature_usage( + "anthropic_federation", + "token_exchange", + if exchanged.is_ok() { "ok" } else { "failed" }, + ); + exchanged +} + +async fn request_federated_token( + client: &reqwest::Client, + base_url: &str, + federation: &AnthropicFederation, + assertion: &str, +) -> Result { + let url = AIProvider::build_anthropic_api_url(base_url, "oauth/token"); + let response = client + .post(url) + .header("anthropic-version", ANTHROPIC_VERSION_STANDARD) + .json(&FederatedTokenRequest { + grant_type: "urn:ietf:params:oauth:grant-type:jwt-bearer", + assertion, + federation_rule_id: &federation.federation_rule_id, + organization_id: &federation.organization_id, + service_account_id: &federation.service_account_id, + workspace_id: federation.workspace_id.as_deref(), + }) + .send() + .await + .map_err(|e| Error::internal_err(format!("Anthropic token exchange failed: {e}")))?; + + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + // Anthropic answers every denied assertion with the same opaque 401; the + // reason is only recorded in the Console's authentication history. + let hint = if status == reqwest::StatusCode::UNAUTHORIZED { + " The reason the exchange was denied is in the authentication history of the Claude Console's Workload identity page." + } else { + "" + }; + return Err(Error::BadRequest(format!( + "Anthropic token exchange failed with status {status}: {body}.{hint}" + ))); + } + + let token = response + .json::() + .await + .map_err(|e| { + Error::internal_err(format!( + "Failed to parse the Anthropic token exchange response: {e}" + )) + })?; + Ok(FederatedAccessToken::new( + token.access_token, + std::time::Duration::from_secs(token.expires_in), + )) +} + +/// The access token of a federated resource for a job, reusing `cached` while it +/// is still fresh, so the agent loop exchanges once per job rather than per turn. +/// Returns `None` for a resource that does not federate. +pub async fn refresh_federated_access_token<'a>( + credentials: &crate::credentials::ProviderCredentials, + cached: &'a mut Option, + client: &AuthedClient, +) -> Result, Error> { + let Some(federation) = credentials.anthropic_federation.as_ref() else { + return Ok(None); + }; + if !cached.as_ref().is_some_and(FederatedAccessToken::is_fresh) { + // A worker holds no OIDC signing key, so it asks the API to mint the token for this job. + let assertion = client + .get_id_token_expiring_in( + ANTHROPIC_OIDC_AUDIENCE, + Some(ANTHROPIC_OIDC_TOKEN_LIFETIME_SECS), + ) + .await + .map_err(|e| Error::ExecutionErr(format!("Failed to get an OIDC token: {e}")))?; + let http_client = crate::utils::pinned_ai_client_for(&credentials.base_url).await?; + *cached = Some( + exchange_federated_token(&http_client, &credentials.base_url, federation, &assertion) + .await?, + ); + } + Ok(cached.as_ref().map(|token| token.access_token.as_str())) +} + +/// Headers authenticating a request with a federated access token, which Anthropic +/// takes as a bearer token in place of `x-api-key`. +pub fn federated_auth_headers(access_token: &str) -> Vec<(&'static str, String)> { + vec![ + ("Authorization", format!("Bearer {access_token}")), + ("anthropic-version", ANTHROPIC_VERSION_STANDARD.to_string()), + ] +} + #[cfg(test)] mod tests { use super::*; @@ -859,6 +1023,7 @@ mod tests { aws_secret_access_key: None, aws_session_token: None, oidc_role_arn: None, + anthropic_federation: None, platform, custom_headers: HashMap::new(), } diff --git a/backend/windmill-ai/src/providers/google_ai.rs b/backend/windmill-ai/src/providers/google_ai.rs index bbf45e5e57..06a404e5a4 100644 --- a/backend/windmill-ai/src/providers/google_ai.rs +++ b/backend/windmill-ai/src/providers/google_ai.rs @@ -805,6 +805,7 @@ mod tests { aws_secret_access_key: None, aws_session_token: None, oidc_role_arn: None, + anthropic_federation: None, platform, custom_headers: HashMap::new(), } diff --git a/backend/windmill-ai/src/providers/mod.rs b/backend/windmill-ai/src/providers/mod.rs index 82dceaa02f..ffddd33d3a 100644 --- a/backend/windmill-ai/src/providers/mod.rs +++ b/backend/windmill-ai/src/providers/mod.rs @@ -370,6 +370,7 @@ mod parity_tests { aws_secret_access_key: None, aws_session_token: None, oidc_role_arn: None, + anthropic_federation: None, platform: kase.platform.clone(), custom_headers: HashMap::new(), } diff --git a/backend/windmill-ai/src/proxy.rs b/backend/windmill-ai/src/proxy.rs index ca6ee1b59b..66285437e4 100644 --- a/backend/windmill-ai/src/proxy.rs +++ b/backend/windmill-ai/src/proxy.rs @@ -285,6 +285,7 @@ mod tests { aws_secret_access_key: None, aws_session_token: None, oidc_role_arn: None, + anthropic_federation: None, platform: AIPlatform::Standard, custom_headers: HashMap::new(), } diff --git a/backend/windmill-ai/src/types.rs b/backend/windmill-ai/src/types.rs index 59733e352c..3950576aed 100644 --- a/backend/windmill-ai/src/types.rs +++ b/backend/windmill-ai/src/types.rs @@ -19,7 +19,7 @@ pub struct McpToolSource { use crate::{ ai_google::sanitize_schema_for_google, ai_providers::{empty_string_as_none, AIProvider}, - credentials::ProviderCredentials, + credentials::{AnthropicFederation, ProviderCredentials}, }; use windmill_common::{db::DB, error::Error, flow_status::AgentAction, flows::FlowModule}; use windmill_parser::Typ; @@ -270,6 +270,16 @@ pub struct ProviderResource { /// credentials exchange this crate does not perform. #[serde(default, deserialize_with = "empty_string_as_none")] pub token_url: Option, + /// Set only by `anthropic_federated` resources, which authenticate through + /// Workload Identity Federation with an OIDC token minted for the job. + #[serde(default, deserialize_with = "empty_string_as_none")] + pub federation_rule_id: Option, + #[serde(default, deserialize_with = "empty_string_as_none")] + pub organization_id: Option, + #[serde(default, deserialize_with = "empty_string_as_none")] + pub service_account_id: Option, + #[serde(default, deserialize_with = "empty_string_as_none")] + pub workspace_id: Option, /// The API token of a `cloudflare` resource, which names it `token` rather than `api_key`. #[serde(default, deserialize_with = "empty_string_as_none")] pub token: Option, @@ -323,8 +333,13 @@ impl ProviderWithResource { /// Callers must only pass resources that were already authorized for the /// current job/workspace; this helper does not perform access checks. pub async fn to_provider_credentials(&self, db: &DB) -> Result { + let anthropic_federation = self.anthropic_federation()?; let base_url = if self.kind == AIProvider::AWSBedrock { String::new() + } else if anthropic_federation.is_some() { + // The exchanged token is the job's identity, so it only ever goes to + // Anthropic: a resource cannot point it at another host. + crate::providers::anthropic::ANTHROPIC_API_BASE_URL.to_string() } else { self.get_base_url(db).await? }; @@ -348,11 +363,40 @@ impl ProviderWithResource { aws_secret_access_key: self.resource.aws_secret_access_key.clone(), aws_session_token: self.resource.aws_session_token.clone(), oidc_role_arn: self.resource.oidc_role_arn.clone(), + anthropic_federation, platform: self.resource.platform.clone(), custom_headers: self.resource.headers.clone(), }) } + /// The federation settings of an `anthropic_federated` resource, which only + /// `federation_rule_id` identifies; the other required fields are checked here. + fn anthropic_federation(&self) -> Result, Error> { + let resource = &self.resource; + let Some(federation_rule_id) = resource.federation_rule_id.clone() else { + return Ok(None); + }; + if self.kind != AIProvider::Anthropic { + return Err(Error::ExecutionErr(format!( + "An anthropic_federated resource can only be used with the Anthropic provider, not {:?}", + self.kind + ))); + } + let required = |value: &Option, name: &str| { + value.clone().ok_or_else(|| { + Error::ExecutionErr(format!( + "The anthropic_federated resource needs {name} to be set" + )) + }) + }; + Ok(Some(AnthropicFederation { + federation_rule_id, + organization_id: required(&resource.organization_id, "organization_id")?, + service_account_id: required(&resource.service_account_id, "service_account_id")?, + workspace_id: resource.workspace_id.clone(), + })) + } + #[cfg(feature = "bedrock")] pub fn get_region(&self) -> Option<&str> { self.resource.region.as_deref() @@ -983,6 +1027,47 @@ pub struct S3ObjectWithType { mod tests { use super::*; + #[test] + fn a_federated_resource_resolves_its_federation_and_names_a_missing_field() { + let provider = |kind: &str, resource: &str| -> ProviderWithResource { + serde_json::from_str(&format!( + r#"{{"kind":"{kind}","model":"m","resource":{resource}}}"# + )) + .unwrap() + }; + let federated = r#"{"federation_rule_id":"fdrl_x","organization_id":"org","service_account_id":"svac_x","workspace_id":""}"#; + + assert_eq!( + provider("anthropic", federated) + .anthropic_federation() + .unwrap(), + Some(AnthropicFederation { + federation_rule_id: "fdrl_x".to_string(), + organization_id: "org".to_string(), + service_account_id: "svac_x".to_string(), + workspace_id: None, + }) + ); + let incomplete = provider( + "anthropic", + r#"{"federation_rule_id":"fdrl_x","organization_id":"org"}"#, + ); + assert!(incomplete + .anthropic_federation() + .unwrap_err() + .to_string() + .contains("service_account_id")); + assert!(provider("openai", federated) + .anthropic_federation() + .is_err()); + assert_eq!( + provider("anthropic", r#"{"api_key":"k","organization_id":"org"}"#) + .anthropic_federation() + .unwrap(), + None + ); + } + /// The schema is part of the prompt prefix the provider caches on, so its keys must /// serialize in the same order on every request. #[test] diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index 8de5e16c11..213a41b4eb 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -24,7 +24,10 @@ use windmill_ai::ai_providers::{ empty_string_as_none, AIPlatform, AIProvider, ProviderConfig, ProviderModel, }; use windmill_ai::ai_types::{validate_token_limit, CONTEXT_WINDOWS, MAX_MODEL_RATE, OUTPUT_LIMITS}; -use windmill_ai::credentials::ProviderCredentials; +use windmill_ai::credentials::{AnthropicFederation, ProviderCredentials}; +use windmill_ai::providers::anthropic::{ + exchange_federated_token, FederatedAccessToken, ANTHROPIC_API_BASE_URL, +}; #[cfg(feature = "bedrock")] use windmill_ai::providers::bedrock::{ handle_bedrock_proxy, BedrockProxyResponse, BedrockProxyResponseBody, @@ -119,6 +122,10 @@ lazy_static::lazy_static! { pub static ref AI_REQUEST_CACHE: Cache<(String, AIProvider), ExpiringProviderCredentials> = Cache::new(500); + /// Anthropic access tokens from a federated exchange, keyed by (workspace, + /// federation settings, user email). Like BEDROCK_ASSUMED_ROLE_CACHE, an entry + /// is never handed to another user than the one the OIDC token named. + static ref ANTHROPIC_FEDERATED_TOKEN_CACHE: Cache<(String, AnthropicFederation, String), FederatedAccessToken> = Cache::new(500); } #[cfg(feature = "bedrock")] @@ -132,6 +139,8 @@ lazy_static::lazy_static! { pub(crate) fn invalidate_ai_request_cache_for_workspace(workspace_id: &str) { AI_REQUEST_CACHE.retain(|(cached_workspace_id, _), _| cached_workspace_id != workspace_id); + ANTHROPIC_FEDERATED_TOKEN_CACHE + .retain(|(cached_workspace_id, _, _), _| cached_workspace_id != workspace_id); // Dropped alongside the credentials that named the role, so that changing a // workspace's AI settings does not leave a session assumed under the old ones // in play until STS expires it. @@ -247,10 +256,33 @@ struct OAuthTokens { access_token: String, } +/// An `anthropic_federated` resource: the Claude Console federation rule, +/// organization, service account and (for a rule enabled in several workspaces) +/// workspace to exchange the requesting user's OIDC token against. +/// +/// It has no base URL on purpose. The exchanged token is the identity of whoever +/// uses the resource, not of whoever wrote it, so a configurable host would let +/// the resource's author collect other users' tokens. +/// +/// Only `federation_rule_id` decides that a value is one: the other fields are +/// optional here and checked afterwards, so a resource missing one is reported +/// rather than read as a standard resource with no key. +#[derive(Deserialize, Debug)] +struct AIAnthropicFederatedResource { + federation_rule_id: String, + #[serde(default)] + organization_id: String, + #[serde(default)] + service_account_id: String, + #[serde(default, deserialize_with = "empty_string_as_none")] + workspace_id: Option, +} + #[derive(Deserialize, Debug)] #[serde(untagged)] enum AIResource { OAuth(AIOAuthResource), + AnthropicFederated(AIAnthropicFederatedResource), Standard(AIStandardResource), } @@ -345,10 +377,50 @@ async fn resolve_provider_credentials( aws_secret_access_key, aws_session_token, oidc_role_arn, + anthropic_federation: None, platform: resource.platform, custom_headers: resource.headers, }) } + AIResource::AnthropicFederated(resource) => { + if *provider != AIProvider::Anthropic { + return Err(Error::BadRequest(format!( + "An anthropic_federated resource can only be used for the Anthropic provider, not {provider:?}" + ))); + } + let required = [ + ("federation_rule_id", &resource.federation_rule_id), + ("organization_id", &resource.organization_id), + ("service_account_id", &resource.service_account_id), + ]; + if let Some((field, _)) = required.iter().find(|(_, value)| value.is_empty()) { + return Err(Error::BadRequest(format!( + "The anthropic_federated resource needs {field} to be set" + ))); + } + + Ok(ProviderCredentials { + provider: provider.clone(), + base_url: ANTHROPIC_API_BASE_URL.to_string(), + api_key: None, + access_token: None, + organization_id: None, + user: None, + region: None, + aws_access_key_id: None, + aws_secret_access_key: None, + aws_session_token: None, + oidc_role_arn: None, + anthropic_federation: Some(AnthropicFederation { + federation_rule_id: resource.federation_rule_id, + organization_id: resource.organization_id, + service_account_id: resource.service_account_id, + workspace_id: resource.workspace_id, + }), + platform: AIPlatform::Standard, + custom_headers: HashMap::new(), + }) + } AIResource::OAuth(resource) => { let user = if let Some(user) = resource.user.clone() { Some(resolve_var(user, db, w_id, user_db.as_ref(), authed).await?) @@ -370,6 +442,7 @@ async fn resolve_provider_credentials( aws_secret_access_key: None, aws_session_token: None, oidc_role_arn: None, + anthropic_federation: None, platform: AIPlatform::Standard, custom_headers: HashMap::new(), }) @@ -522,6 +595,73 @@ async fn assume_bedrock_role( } } +/// Exchange an OIDC token minted for the requesting user for an Anthropic access +/// token, in place, when the resource is set up for Workload Identity Federation. +/// +/// Runs after the workspace credential cache for the same reason the Bedrock +/// exchange does: the token's `sub` names the user, and is what a federation rule +/// matches and the Claude Console's authentication history records. +async fn resolve_anthropic_federated_credentials( + db: &DB, + w_id: &str, + authed: &ApiAuthed, + credentials: &mut ProviderCredentials, +) -> Result<()> { + let Some(federation) = credentials.anthropic_federation.as_ref() else { + return Ok(()); + }; + + let cache_key = (w_id.to_string(), federation.clone(), authed.email.clone()); + let token = match ANTHROPIC_FEDERATED_TOKEN_CACHE + .get(&cache_key) + .filter(FederatedAccessToken::is_fresh) + { + Some(cached) => cached, + None => { + let assertion = mint_anthropic_oidc_token(db, w_id, authed).await?; + let client = pinned_ai_client_for(&credentials.base_url).await?; + let token = + exchange_federated_token(&client, &credentials.base_url, federation, &assertion) + .await?; + ANTHROPIC_FEDERATED_TOKEN_CACHE.insert(cache_key, token.clone()); + token + } + }; + + credentials.access_token = Some(token.access_token); + Ok(()) +} + +async fn mint_anthropic_oidc_token(db: &DB, w_id: &str, authed: &ApiAuthed) -> Result { + #[cfg(all(feature = "enterprise", feature = "openidconnect", feature = "private"))] + { + use windmill_ai::providers::anthropic::{ + ANTHROPIC_OIDC_AUDIENCE, ANTHROPIC_OIDC_TOKEN_LIFETIME_SECS, + }; + use windmill_common::oidc_oss::{generate_id_token, OidcCaller, WorkspaceClaim}; + + let id_token = generate_id_token( + OidcCaller::Anthropic, + Some(db), + WorkspaceClaim { workspace: w_id.to_string() }, + ANTHROPIC_OIDC_AUDIENCE, + format!("{}::{}", authed.email, w_id), + Some(authed.email.clone()), + Some(ANTHROPIC_OIDC_TOKEN_LIFETIME_SECS), + ) + .await?; + Ok(id_token.to_string()) + } + #[cfg(not(all(feature = "enterprise", feature = "openidconnect", feature = "private")))] + { + let _ = (db, w_id, authed); + Err(Error::BadRequest( + "Anthropic Workload Identity Federation through Windmill's OIDC provider requires an enterprise license" + .to_string(), + )) + } +} + #[derive(Clone, Debug)] pub struct ExpiringProviderCredentials { credentials: ProviderCredentials, @@ -1129,6 +1269,7 @@ async fn global_proxy( aws_secret_access_key: None, aws_session_token: None, oidc_role_arn: None, + anthropic_federation: None, platform: AIPlatform::Standard, custom_headers: HashMap::new(), }; @@ -1435,6 +1576,7 @@ async fn proxy( #[cfg(feature = "bedrock")] resolve_bedrock_oidc_credentials(&db, &w_id, &authed, &mut credentials).await?; + resolve_anthropic_federated_credentials(&db, &w_id, &authed, &mut credentials).await?; // Free tier: pin the model and clamp max_tokens server-side before forwarding, // since the request body is otherwise client-controlled. @@ -1643,11 +1785,41 @@ mod tests { aws_secret_access_key: None, aws_session_token: None, oidc_role_arn: None, + anthropic_federation: None, platform: AIPlatform::Standard, custom_headers: HashMap::new(), } } + #[test] + fn only_a_federation_rule_makes_a_resource_anthropic_federated() { + let federated = serde_json::json!({ + "federation_rule_id": "fdrl_x", + "organization_id": "00000000-0000-0000-0000-000000000000", + "service_account_id": "svac_x", + "workspace_id": "" + }); + assert!(matches!( + serde_json::from_value::(federated), + Ok(AIResource::AnthropicFederated( + AIAnthropicFederatedResource { workspace_id: None, .. } + )) + )); + + let incomplete = serde_json::json!({"federation_rule_id": "fdrl_x"}); + assert!(matches!( + serde_json::from_value::(incomplete), + Ok(AIResource::AnthropicFederated(_)) + )); + + // `organization_id` is also the OpenAI organization of a standard resource. + let standard = serde_json::json!({"api_key": "k", "organization_id": "org-x"}); + assert!(matches!( + serde_json::from_value::(standard), + Ok(AIResource::Standard(_)) + )); + } + #[test] fn ai_standard_resource_ignores_legacy_enable_1m_context_keys() { // Resources created before the field was removed still carry the legacy key diff --git a/backend/windmill-common/src/client.rs b/backend/windmill-common/src/client.rs index 8edcdc64fe..54eb486da7 100644 --- a/backend/windmill-common/src/client.rs +++ b/backend/windmill-common/src/client.rs @@ -71,15 +71,32 @@ impl AuthedClient { } pub async fn get_id_token(&self, audience: &str) -> anyhow::Result { + self.get_id_token_expiring_in(audience, None).await + } + + /// `expires_in` caps the token's lifetime below the instance default, for + /// audiences that refuse long-lived assertions. + pub async fn get_id_token_expiring_in( + &self, + audience: &str, + expires_in: Option, + ) -> anyhow::Result { + // An audience can be a URL, whose slashes would otherwise split the path segment. let url = format!( "{}/api/w/{}/oidc/token/{}", - self.base_internal_url, self.workspace, audience + self.base_internal_url, + self.workspace, + urlencoding::encode(audience) ); - let response = self + let mut request = self .force_client .as_ref() .unwrap_or(&HTTP_CLIENT) - .post(&url) + .post(&url); + if let Some(expires_in) = expires_in { + request = request.query(&[("expires_in", expires_in)]); + } + let response = request .header( reqwest::header::AUTHORIZATION, reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.token))?, diff --git a/backend/windmill-common/src/oidc_oss.rs b/backend/windmill-common/src/oidc_oss.rs index 2945f8b87a..60ef4da687 100644 --- a/backend/windmill-common/src/oidc_oss.rs +++ b/backend/windmill-common/src/oidc_oss.rs @@ -98,6 +98,7 @@ pub enum OidcCaller { Vault, Sqs, Bedrock, + Anthropic, } #[cfg(not(feature = "private"))] diff --git a/backend/windmill-worker/src/ai/compaction.rs b/backend/windmill-worker/src/ai/compaction.rs index b933348a2e..82a79985e8 100644 --- a/backend/windmill-worker/src/ai/compaction.rs +++ b/backend/windmill-worker/src/ai/compaction.rs @@ -610,12 +610,24 @@ async fn summarize_prefix( request .query_builder .get_endpoint(base_url, request.model, &OutputType::Text); - let auth_headers = retain_effective_credentials( + // Summaries are rare enough that exchanging afresh costs less than threading + // the agent loop's cached token through. + let mut federated_token = None; + let auth_headers = match windmill_ai::providers::anthropic::refresh_federated_access_token( request.credentials, - request - .query_builder - .get_auth_headers(api_key, base_url, &OutputType::Text), - ); + &mut federated_token, + request.client, + ) + .await? + { + Some(token) => windmill_ai::providers::anthropic::federated_auth_headers(token), + None => retain_effective_credentials( + request.credentials, + request + .query_builder + .get_auth_headers(api_key, base_url, &OutputType::Text), + ), + }; // As in the agent loop, `endpoint` derives from the user-controlled provider // base_url, so DNS is pinned to the SSRF-validated address. diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index bf5c3e4f9f..6c37f67f2e 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1666,6 +1666,9 @@ pub async fn run_agent( // re-assumed only near expiry, rather than on every model turn. #[cfg(feature = "bedrock")] let mut bedrock_assumed_role: Option = None; + let mut anthropic_federated_token: Option< + windmill_ai::providers::anthropic::FederatedAccessToken, + > = None; // Main agent loop for i in 0..max_iterations { @@ -1804,6 +1807,15 @@ pub async fn run_agent( ))); } + let federated_token = + windmill_ai::providers::anthropic::refresh_federated_access_token( + &credentials, + &mut anthropic_federated_token, + client, + ) + .await? + .map(str::to_string); + let timeout = resolve_job_timeout(conn, &job.workspace_id, job.id, job.timeout) .await .0; @@ -1855,10 +1867,15 @@ pub async fn run_agent( }; let endpoint = query_builder.get_endpoint(base_url, args.provider.get_model(), output_type); - let auth_headers = retain_effective_credentials( - &credentials, - query_builder.get_auth_headers(api_key, base_url, output_type), - ); + let auth_headers = match federated_token.as_deref() { + Some(token) => { + windmill_ai::providers::anthropic::federated_auth_headers(token) + } + None => retain_effective_credentials( + &credentials, + query_builder.get_auth_headers(api_key, base_url, output_type), + ), + }; let resp = build_http_request(&endpoint, &auth_headers, request_body) .send() diff --git a/frontend/src/lib/components/AIProviderPicker.svelte b/frontend/src/lib/components/AIProviderPicker.svelte index 67afad9be6..0d23ab55a3 100644 --- a/frontend/src/lib/components/AIProviderPicker.svelte +++ b/frontend/src/lib/components/AIProviderPicker.svelte @@ -9,6 +9,7 @@ } from './copilot/lib' import type { AIProvider, ProviderConfig } from '$lib/gen' import ResourcePicker from './ResourcePicker.svelte' + import { aiResourceProvider } from './copilot/aiResourceProvider' import Toggle from './Toggle.svelte' import { saveConfig, removeConfig, isSameAsStoredConfig } from './aiProviderStorage' import AIReasoningEffortPicker from './AIReasoningEffortPicker.svelte' @@ -57,7 +58,11 @@ const offeredProviders = Object.keys( untrack(() => decision) ? DECISION_AI_PROVIDERS : AI_PROVIDERS ) - const providerResourceTypes = offeredProviders.join(',') + // Federated Anthropic resources authenticate the anthropic provider without a key of their own. + const providerResourceTypes = [ + ...offeredProviders, + ...(offeredProviders.includes('anthropic') ? ['anthropic_federated'] : []) + ].join(',') if (!_uncheckedValue) { _uncheckedValue = { @@ -143,13 +148,14 @@ function onResourcePicked(_path: string | undefined, type: string | undefined) { // An empty type is the placeholder the picker keeps for a saved path it could not find. It // says nothing about the provider, so the kind stands. - if (!value || !type || !offeredProviders.includes(type)) { + const provider = aiResourceProvider(type) + if (!value || !provider || !offeredProviders.includes(provider)) { return } - if (value.kind === type) { + if (value.kind === provider) { return } - value.kind = type as AIProvider + value.kind = provider // Models are per provider, and a reasoning token is per model. value.model = '' value.reasoning_effort = undefined diff --git a/frontend/src/lib/components/AppConnectInner.svelte b/frontend/src/lib/components/AppConnectInner.svelte index 81dd5225fb..069965b4e7 100644 --- a/frontend/src/lib/components/AppConnectInner.svelte +++ b/frontend/src/lib/components/AppConnectInner.svelte @@ -54,6 +54,7 @@ import SchemaForm from './SchemaForm.svelte' import { useOperatingWorkspace } from '$lib/components/operatingWorkspace.svelte' import type { NativeConnect } from './triggers/native/utils' + import { AI_RESOURCE_TYPE_ALTERNATIVES } from './copilot/aiResourceProvider' const operatingWorkspace = useOperatingWorkspace() @@ -258,6 +259,16 @@ return registryEntryFor(connectClient, resourceType) } + let resourceTypeAlternative = $derived(AI_RESOURCE_TYPE_ALTERNATIVES[resourceType]) + + function switchToResourceType(rt: string) { + connectClient = rt + resourceType = rt + linkedSecrets = [] + args = {} + getResourceTypeInfo() + } + /** The static registry declares this provider supports client credentials */ function registryCcCapable(): boolean { return registryCcCapableFor(connectClient, resourceType) @@ -1477,6 +1488,14 @@ Acquire the token automatically via client credentials instead {/if} + {#if resourceTypeAlternative} + + {/if}
diff --git a/frontend/src/lib/components/ResourcePicker.svelte b/frontend/src/lib/components/ResourcePicker.svelte index b396bd6dee..d766e5a51f 100644 --- a/frontend/src/lib/components/ResourcePicker.svelte +++ b/frontend/src/lib/components/ResourcePicker.svelte @@ -20,6 +20,9 @@ value?: string | undefined valueType?: string | undefined resourceType?: string | undefined + /** Further resource types whose resources are offered alongside `resourceType`'s. New + * resources are still created as `resourceType`. */ + additionalResourceTypes?: string[] disabled?: boolean disablePortal?: boolean showSchemaExplorer?: boolean @@ -46,6 +49,7 @@ value = $bindable(undefined), valueType = $bindable(undefined), resourceType = undefined, + additionalResourceTypes = [], disabled = false, disablePortal = false, showSchemaExplorer = false, @@ -114,8 +118,10 @@ async function loadResources(resourceType: string | undefined) { loading = true try { - const resourceTypesToQuery = - resourceType === 'snowflake' ? ['snowflake', 'snowflake_oauth'] : [resourceType] + const resourceTypesToQuery = [ + ...(resourceType === 'snowflake' ? ['snowflake', 'snowflake_oauth'] : [resourceType]), + ...additionalResourceTypes + ] const resources = await Promise.all( resourceTypesToQuery.map((rt) => diff --git a/frontend/src/lib/components/copilot/aiResourceProvider.ts b/frontend/src/lib/components/copilot/aiResourceProvider.ts index 91ae926dfc..cd80bb21c0 100644 --- a/frontend/src/lib/components/copilot/aiResourceProvider.ts +++ b/frontend/src/lib/components/copilot/aiResourceProvider.ts @@ -1,9 +1,20 @@ import type { AIProvider } from '$lib/gen' import { AI_PROVIDERS } from './lib' +/** A resource type that authenticates the same provider another way, offered as a switch when + * creating a resource of the first, since a picker scoped to a provider only creates its main type. */ +export const AI_RESOURCE_TYPE_ALTERNATIVES: Record = { + anthropic: { + resourceType: 'anthropic_federated', + label: 'Authenticate with Workload Identity Federation instead of an API key' + }, + anthropic_federated: { resourceType: 'anthropic', label: 'Use an API key instead' } +} + /** The AI provider a resource of this type holds credentials for, if any. */ export function aiResourceProvider(resourceType: string | undefined): AIProvider | undefined { if (resourceType === 'openai_client_credentials_oauth') return 'openai' + if (resourceType === 'anthropic_federated') return 'anthropic' return resourceType && Object.hasOwn(AI_PROVIDERS, resourceType) ? (resourceType as AIProvider) : undefined diff --git a/frontend/src/lib/components/workspaceSettings/AISettings.svelte b/frontend/src/lib/components/workspaceSettings/AISettings.svelte index 05159e3057..79d7be2214 100644 --- a/frontend/src/lib/components/workspaceSettings/AISettings.svelte +++ b/frontend/src/lib/components/workspaceSettings/AISettings.svelte @@ -520,6 +520,9 @@ resourceType={provider === 'openai' && usingOpenaiClientCredentialsOauth ? 'openai_client_credentials_oauth' : provider} + additionalResourceTypes={provider === 'anthropic' + ? ['anthropic_federated'] + : []} initialValue={aiProviders[provider].resource_path} bind:value={ () => aiProviders[provider].resource_path || undefined,