From e14da5c6bc8256bfe1d71ddecb886c3efc994ae4 Mon Sep 17 00:00:00 2001 From: hugocasa Date: Fri, 25 Sep 2026 18:19:38 +0200 Subject: [PATCH] feat(bedrock): add OIDC role assumption as a fourth auth mode (#10936) * feat(bedrock): add OIDC role assumption as a fourth auth mode Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01YCVb91fZp3dqM14KRPzoEn * fix(bedrock): gate the OIDC cache correctly and assume the role once per job Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01YCVb91fZp3dqM14KRPzoEn * refactor(bedrock): check the OIDC region before minting a token Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01YCVb91fZp3dqM14KRPzoEn * fix(bedrock): keep OIDC session names collision-resistant, gate the copy on EE Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01YCVb91fZp3dqM14KRPzoEn * fix(bedrock): check the OIDC region before reusing cached credentials Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01YCVb91fZp3dqM14KRPzoEn * fix(bedrock): clear assumed-role sessions when AI settings change invalidate_ai_request_cache_for_workspace cleared AI_REQUEST_CACHE only, so a workspace's AI settings edit reset one cache and left the assumed-role sessions keyed on the old config in place until STS expired them. Also name the region requirement in the credentials-check hint, so following it does not land on the OIDC path's region guard. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01YCVb91fZp3dqM14KRPzoEn * chore: update ee-repo-ref to de73db2bacfdc3eaa2e63b1827178bc198d54e5c This commit updates the EE repository reference after PR #770 was merged in windmill-ee-private. Previous ee-repo-ref: c43dab1e69b1cb3f685e6df07bff634dc2a0b734 New ee-repo-ref: de73db2bacfdc3eaa2e63b1827178bc198d54e5c Automated by sync-ee-ref workflow. --------- Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: windmill-internal-app[bot] --- backend/ee-repo-ref.txt | 2 +- backend/windmill-ai/Cargo.toml | 2 +- backend/windmill-ai/src/ai_bedrock.rs | 258 ++++++++++++++++++ backend/windmill-ai/src/credentials.rs | 13 + .../windmill-ai/src/providers/anthropic.rs | 1 + .../windmill-ai/src/providers/google_ai.rs | 1 + backend/windmill-ai/src/providers/mod.rs | 1 + backend/windmill-ai/src/proxy.rs | 1 + backend/windmill-ai/src/types.rs | 7 + backend/windmill-api/src/ai.rs | 156 +++++++++++ backend/windmill-worker/src/ai_executor.rs | 32 ++- .../components/BedrockCredentialsCheck.svelte | 6 + 12 files changed, 475 insertions(+), 5 deletions(-) diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 06b1b10a9a..f102167077 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -c099af43bad922fd57ed0449ab717ffa047b5546 +de73db2bacfdc3eaa2e63b1827178bc198d54e5c diff --git a/backend/windmill-ai/Cargo.toml b/backend/windmill-ai/Cargo.toml index 98ea87a4e3..b1dea3ae54 100644 --- a/backend/windmill-ai/Cargo.toml +++ b/backend/windmill-ai/Cargo.toml @@ -6,7 +6,7 @@ edition.workspace = true [features] default = [] -bedrock = ["dep:aws-sdk-bedrock", "dep:aws-sdk-bedrockruntime", "dep:aws-credential-types", "dep:aws-smithy-types", "dep:aws-config"] +bedrock = ["dep:aws-sdk-bedrock", "dep:aws-sdk-bedrockruntime", "dep:aws-credential-types", "dep:aws-smithy-types", "dep:aws-config", "windmill-common/aws_auth"] mcp = ["dep:windmill-mcp"] [lib] diff --git a/backend/windmill-ai/src/ai_bedrock.rs b/backend/windmill-ai/src/ai_bedrock.rs index b223725af2..a9318a1f3a 100644 --- a/backend/windmill-ai/src/ai_bedrock.rs +++ b/backend/windmill-ai/src/ai_bedrock.rs @@ -2,6 +2,7 @@ //! //! This module provides: //! - BedrockClient: SDK wrapper with bearer token and IAM credentials auth +//! - OIDC role assumption: exchange a Windmill OIDC token for temporary IAM keys //! - Message/tool conversion: OpenAI format <-> Bedrock Converse API format //! - Stream event parsing: Extract text/tool deltas from Bedrock stream events //! @@ -291,6 +292,187 @@ impl BedrockClient { } } +// ============================================================================ +// OIDC Role Assumption +// ============================================================================ + +pub use windmill_common::auth::aws::AWS_OIDC_AUDIENCE; + +/// Bedrock credential precedence: a bearer API key wins, then a complete pair of +/// explicit IAM keys, then OIDC role assumption, then the ambient environment. +/// Returns the role to assume only when both higher-priority modes are unset, so +/// a resource can carry a role ARN without it ever overriding keys set on it. +pub fn bedrock_oidc_role_to_assume<'a>( + api_key: Option<&str>, + aws_access_key_id: Option<&str>, + aws_secret_access_key: Option<&str>, + oidc_role_arn: Option<&'a str>, +) -> Option<&'a str> { + let has_bearer_token = api_key.is_some_and(|key| !key.is_empty()); + let has_iam_keys = aws_access_key_id.is_some_and(|key| !key.is_empty()) + && aws_secret_access_key.is_some_and(|key| !key.is_empty()); + + if has_bearer_token || has_iam_keys { + return None; + } + + oidc_role_arn.filter(|arn| !arn.is_empty()) +} + +/// Re-assume the role this long before AWS expires the credentials, so a request +/// that starts on a reused set cannot finish holding dead ones. +pub const ASSUMED_ROLE_REFRESH_MARGIN: std::time::Duration = std::time::Duration::from_secs(120); + +/// Temporary IAM credentials minted by STS `AssumeRoleWithWebIdentity`. +#[derive(Clone, Debug)] +pub struct AssumedRoleCredentials { + pub access_key_id: String, + pub secret_access_key: String, + pub session_token: String, + pub expires_at: std::time::SystemTime, +} + +impl AssumedRoleCredentials { + /// Whether these have enough life left to sign another request. + pub fn is_fresh(&self) -> bool { + self.expires_at > std::time::SystemTime::now() + ASSUMED_ROLE_REFRESH_MARGIN + } +} + +/// Assume the resource's OIDC role for a job, reusing `cached` while it is still +/// fresh, and hand back the temporary keys to sign Bedrock requests with. +/// +/// The agent loop calls this once per iteration, so without the reuse a +/// tool-heavy run would mint an OIDC token and call STS on every model turn. +/// Returns `None` when a higher-priority credential is set, so the caller can +/// call it unconditionally. +pub async fn refresh_bedrock_oidc_credentials<'a>( + credentials: &crate::credentials::ProviderCredentials, + cached: &'a mut Option, + client: &windmill_common::client::AuthedClient, + job_id: &uuid::Uuid, +) -> Result, Error> { + let Some(role_arn) = bedrock_oidc_role_to_assume( + credentials.api_key.as_deref(), + credentials.aws_access_key_id.as_deref(), + credentials.aws_secret_access_key.as_deref(), + credentials.oidc_role_arn.as_deref(), + ) else { + return Ok(None); + }; + + // Ahead of the reuse check, so a request never signs against the empty region + // just because an earlier one had already assumed the role. + let region = bedrock_oidc_region(credentials.region.as_deref())?; + + if !cached + .as_ref() + .is_some_and(AssumedRoleCredentials::is_fresh) + { + // A worker holds no OIDC signing key, so it asks the API to mint the + // token for this job; the session name carries the job into CloudTrail. + let id_token = client + .get_id_token(AWS_OIDC_AUDIENCE) + .await + .map_err(|e| Error::internal_err(format!("Failed to get OIDC token: {}", e)))?; + // The whole UUID: "windmill-ai-job-" plus 32 hex is 48 characters, well + // inside AWS's limit, and a truncated one would make two runs + // indistinguishable in CloudTrail. + let session_name = aws_role_session_name("windmill-ai-job", &job_id.simple().to_string()); + *cached = + Some(assume_role_with_oidc_token(role_arn, region, id_token, &session_name).await?); + } + + Ok(cached.as_ref()) +} + +/// Distinguishing digest appended to a session name that had to be truncated. +const SESSION_NAME_DIGEST_LEN: usize = 8; + +/// Build a role session name AWS accepts: it constrains them to +/// `[\w+=,.@-]{2,64}`, and rejects the whole request otherwise. Disallowed +/// characters are replaced rather than dropped. +/// +/// The name is the only thing carrying the caller into the assumed-role ARN and +/// CloudTrail, so an identity too long to fit keeps a prefix plus a digest of the +/// whole thing: two identities sharing a prefix would otherwise be +/// indistinguishable in exactly the audit trail this feature exists to produce. +pub fn aws_role_session_name(prefix: &str, identity: &str) -> String { + let sanitize = |s: &str| -> String { + s.chars() + .map(|c| { + if c.is_ascii_alphanumeric() || matches!(c, '_' | '+' | '=' | ',' | '.' | '@' | '-') + { + c + } else { + '-' + } + }) + .collect() + }; + + let name = sanitize(&format!("{}-{}", prefix, identity)); + if name.len() <= 64 { + return name; + } + + let digest = &windmill_common::utils::calculate_hash(identity)[..SESSION_NAME_DIGEST_LEN]; + let kept = 64 - SESSION_NAME_DIGEST_LEN - 1; + format!("{}-{}", &name[..kept], digest) +} + +/// The region an OIDC role assumption runs in. +/// +/// The assumption resolves to explicit IAM keys, and the SDK client built from +/// those has no ambient-region fallback the way [`BedrockClient::from_env`] does, +/// so a resource that assumes a role has to name its region. Callers check this +/// before minting the OIDC token, so a misconfigured resource never issues an +/// identity token it is only going to throw away. +pub fn bedrock_oidc_region(region: Option<&str>) -> Result<&str, Error> { + region.filter(|r| !r.is_empty()).ok_or_else(|| { + Error::BadRequest( + "AWS Bedrock resources that assume a role through OIDC must set a region".to_string(), + ) + }) +} + +/// Exchange a Windmill OIDC token for temporary IAM credentials on `role_arn`. +pub async fn assume_role_with_oidc_token( + role_arn: &str, + region: &str, + id_token: String, + session_name: &str, +) -> Result { + use windmill_common::auth::aws::{ + get_assume_role_with_web_identity_fluent_builder, GetAuthenticationOutput, OidcAuth, + }; + + let oidc_auth = OidcAuth { region: Some(region.to_string()), role_arn: role_arn.to_string() }; + + let output = + get_assume_role_with_web_identity_fluent_builder(&oidc_auth, id_token, Some(session_name)) + .await? + .send() + .await + .map_err(|e| { + Error::internal_err(format!( + "Failed to assume AWS role {} with Windmill's OIDC token: {}", + role_arn, + format_bedrock_error(&e) + )) + })?; + + let credentials = output.get_credentials()?; + + Ok(AssumedRoleCredentials { + access_key_id: credentials.access_key_id.clone(), + secret_access_key: credentials.secret_access_key.clone(), + session_token: credentials.session_token.clone(), + expires_at: std::time::UNIX_EPOCH + + std::time::Duration::from_secs(credentials.expiration.secs().max(0) as u64), + }) +} + // ============================================================================ // Error Formatting // ============================================================================ @@ -1000,6 +1182,82 @@ mod tests { use aws_sdk_bedrockruntime::types::ReasoningContentBlock; use serde_json::value::RawValue; + #[test] + fn oidc_role_yields_only_to_bearer_and_complete_iam_keys() { + let arn = Some("arn:aws:iam::1:role/bedrock"); + + assert_eq!( + bedrock_oidc_role_to_assume(None, None, None, arn), + arn, + "nothing else set: assume the role rather than fall through to the environment" + ); + assert_eq!( + bedrock_oidc_role_to_assume(Some("key"), None, None, arn), + None + ); + assert_eq!( + bedrock_oidc_role_to_assume(None, Some("id"), Some("secret"), arn), + None + ); + assert_eq!( + bedrock_oidc_role_to_assume(None, Some("id"), None, arn), + arn, + "half a key pair is not usable IAM credentials, so the role still wins" + ); + // Cleared resource fields arrive as empty strings, not as absent ones. + assert_eq!( + bedrock_oidc_role_to_assume(Some(""), Some(""), Some(""), arn), + arn + ); + assert_eq!( + bedrock_oidc_role_to_assume(None, None, None, Some("")), + None + ); + } + + #[test] + fn credentials_stop_being_fresh_before_aws_expires_them() { + let at = |d: std::time::Duration| AssumedRoleCredentials { + access_key_id: String::new(), + secret_access_key: String::new(), + session_token: String::new(), + expires_at: std::time::SystemTime::now() + d, + }; + // Reused only while there is still margin left to finish a request. + assert!(at(ASSUMED_ROLE_REFRESH_MARGIN * 2).is_fresh()); + assert!(!at(ASSUMED_ROLE_REFRESH_MARGIN / 2).is_fresh()); + } + + #[test] + fn role_session_name_stays_within_what_aws_accepts() { + // A whole job UUID fits, so CloudTrail names the exact run. + assert_eq!( + aws_role_session_name("windmill-ai-job", "0a1b2c3d4e5f60718293a4b5c6d7e8f9"), + "windmill-ai-job-0a1b2c3d4e5f60718293a4b5c6d7e8f9" + ); + // Usernames and emails routinely carry characters AWS rejects. + assert_eq!( + aws_role_session_name("windmill-ai-copilot", "ada/lovelace (admin)"), + "windmill-ai-copilot-ada-lovelace--admin-" + ); + + // Identities too long to fit keep a digest, so a shared prefix does not + // collapse two callers into one CloudTrail identity. + let a = aws_role_session_name("windmill-ai-copilot", &format!("{}a", "x".repeat(60))); + let b = aws_role_session_name("windmill-ai-copilot", &format!("{}b", "x".repeat(60))); + assert_ne!(a, b); + for name in [&a, &b] { + assert_eq!(name.len(), 64); + assert!(name.chars().all(|c| c.is_ascii_alphanumeric() + || matches!(c, '_' | '+' | '=' | ',' | '.' | '@' | '-'))); + } + + // Sanitising before truncating keeps every char one byte, so the cap is + // never applied mid-character. + let unicode = aws_role_session_name("windmill-ai-copilot", &"é".repeat(100)); + assert_eq!(unicode.len(), 64); + } + fn text_message(role: &str, content: &str) -> OpenAIMessage { OpenAIMessage { role: role.to_string(), diff --git a/backend/windmill-ai/src/credentials.rs b/backend/windmill-ai/src/credentials.rs index c9502bdbb2..5a779938ed 100644 --- a/backend/windmill-ai/src/credentials.rs +++ b/backend/windmill-ai/src/credentials.rs @@ -19,6 +19,19 @@ pub struct ProviderCredentials { pub aws_access_key_id: Option, pub aws_secret_access_key: Option, pub aws_session_token: Option, + /// AWS IAM role to assume through Windmill's OIDC provider. Only consulted + /// when no bearer key and no explicit IAM keys are set — see + /// [`crate::ai_bedrock::bedrock_oidc_role_to_assume`]. The API proxy + /// exchanges it for temporary keys before it builds the request; the worker + /// exchanges it at the Bedrock call, because only the API can sign an OIDC + /// token. + /// + /// The two paths therefore present different tokens to AWS: the proxy a + /// workspace claim with `sub = "::"`, the worker the job + /// claim `/oidc/token` issues, `sub = "::::::"`. + /// 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, pub platform: AIPlatform, pub custom_headers: HashMap, } diff --git a/backend/windmill-ai/src/providers/anthropic.rs b/backend/windmill-ai/src/providers/anthropic.rs index 0bb7316dcb..d274d91fa2 100644 --- a/backend/windmill-ai/src/providers/anthropic.rs +++ b/backend/windmill-ai/src/providers/anthropic.rs @@ -858,6 +858,7 @@ mod tests { aws_access_key_id: None, aws_secret_access_key: None, aws_session_token: None, + oidc_role_arn: 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 8320ef5661..ce37e45df7 100644 --- a/backend/windmill-ai/src/providers/google_ai.rs +++ b/backend/windmill-ai/src/providers/google_ai.rs @@ -780,6 +780,7 @@ mod tests { aws_access_key_id: None, aws_secret_access_key: None, aws_session_token: None, + oidc_role_arn: 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 af9ba3ff9f..908f007978 100644 --- a/backend/windmill-ai/src/providers/mod.rs +++ b/backend/windmill-ai/src/providers/mod.rs @@ -184,6 +184,7 @@ mod parity_tests { aws_access_key_id: None, aws_secret_access_key: None, aws_session_token: None, + oidc_role_arn: 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 c672aa1d42..3af8f88a1d 100644 --- a/backend/windmill-ai/src/proxy.rs +++ b/backend/windmill-ai/src/proxy.rs @@ -263,6 +263,7 @@ mod tests { aws_access_key_id: None, aws_secret_access_key: None, aws_session_token: None, + oidc_role_arn: 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 12e06a35d8..ac65f637e2 100644 --- a/backend/windmill-ai/src/types.rs +++ b/backend/windmill-ai/src/types.rs @@ -245,6 +245,12 @@ pub struct ProviderResource { deserialize_with = "empty_string_as_none" )] pub aws_session_token: Option, + #[serde( + alias = "oidcRoleArn", + default, + deserialize_with = "empty_string_as_none" + )] + pub oidc_role_arn: Option, /// Platform (standard or google_vertex_ai) #[serde(default)] pub platform: AIPlatform, @@ -311,6 +317,7 @@ impl ProviderWithResource { aws_access_key_id: self.resource.aws_access_key_id.clone(), 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(), platform: self.resource.platform.clone(), custom_headers: self.resource.headers.clone(), }) diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index 16b4b17141..3255b629da 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -120,8 +120,23 @@ lazy_static::lazy_static! { } +#[cfg(feature = "bedrock")] +lazy_static::lazy_static! { + /// Temporary IAM credentials from a Bedrock OIDC role assumption, keyed by + /// (workspace, role ARN, user email). The email is the identity the OIDC + /// token is minted for, so unlike AI_REQUEST_CACHE an entry is never handed + /// to a different user — that is what keeps AWS-side attribution per user. + static ref BEDROCK_ASSUMED_ROLE_CACHE: Cache<(String, String, String), windmill_ai::ai_bedrock::AssumedRoleCredentials> = Cache::new(500); +} + pub(crate) fn invalidate_ai_request_cache_for_workspace(workspace_id: &str) { AI_REQUEST_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. + #[cfg(feature = "bedrock")] + BEDROCK_ASSUMED_ROLE_CACHE + .retain(|(cached_workspace_id, _, _), _| cached_workspace_id != workspace_id); } /// Shared configuration for every outbound AI HTTP client (the pooled @@ -219,6 +234,12 @@ struct AIStandardResource { deserialize_with = "empty_string_as_none" )] aws_session_token: Option, + #[serde( + alias = "oidcRoleArn", + default, + deserialize_with = "empty_string_as_none" + )] + oidc_role_arn: Option, /// Platform (standard or google_vertex_ai) #[serde(default)] platform: AIPlatform, @@ -301,6 +322,11 @@ async fn resolve_provider_credentials( } else { None }; + let oidc_role_arn = if let Some(role_arn) = resource.oidc_role_arn { + Some(resolve_var(role_arn, db, w_id, user_db.as_ref(), authed).await?) + } else { + None + }; Ok(ProviderCredentials { provider: provider.clone(), @@ -313,6 +339,7 @@ async fn resolve_provider_credentials( aws_access_key_id, aws_secret_access_key, aws_session_token, + oidc_role_arn, platform: resource.platform, custom_headers: resource.headers, }) @@ -337,6 +364,7 @@ async fn resolve_provider_credentials( aws_access_key_id: None, aws_secret_access_key: None, aws_session_token: None, + oidc_role_arn: None, platform: AIPlatform::Standard, custom_headers: HashMap::new(), }) @@ -385,6 +413,109 @@ async fn get_token_using_oauth( Ok(response.access_token) } +/// Exchange the resource's OIDC role ARN for temporary IAM credentials, in place. +/// +/// Only the API holds the OIDC signing key, so the exchange happens here rather +/// than in the Bedrock handler, which then sees ordinary IAM keys. It runs after +/// the workspace credential cache on purpose: the STS session name names the user +/// making the request, and a cache shared by the whole workspace would hand one +/// user's session to another. +#[cfg(feature = "bedrock")] +async fn resolve_bedrock_oidc_credentials( + db: &DB, + w_id: &str, + authed: &ApiAuthed, + credentials: &mut ProviderCredentials, +) -> Result<()> { + use windmill_ai::ai_bedrock::{aws_role_session_name, bedrock_oidc_role_to_assume}; + + // Only the Bedrock resource type carries a role ARN; leaving the field set on + // any other resource inert keeps a hand-edited one from reaching STS. + if credentials.provider != AIProvider::AWSBedrock { + return Ok(()); + } + + let Some(role_arn) = bedrock_oidc_role_to_assume( + credentials.api_key.as_deref(), + credentials.aws_access_key_id.as_deref(), + credentials.aws_secret_access_key.as_deref(), + credentials.oidc_role_arn.as_deref(), + ) else { + return Ok(()); + }; + + // Ahead of the cache: STS credentials are region-independent, so a resource + // with the same role but no region would otherwise reuse an entry and sign a + // request against the empty region the guard exists to reject. + let region = windmill_ai::ai_bedrock::bedrock_oidc_region(credentials.region.as_deref())?; + + let session_name = aws_role_session_name("windmill-ai-copilot", &authed.username); + // Keyed on the email, the identity the OIDC token is minted for, rather than + // on the session name derived from it: that name is sanitised and length + // capped, so it is a lossy stand-in for the caller. + let cache_key = (w_id.to_string(), role_arn.to_string(), authed.email.clone()); + + let assumed = match BEDROCK_ASSUMED_ROLE_CACHE + .get(&cache_key) + .filter(windmill_ai::ai_bedrock::AssumedRoleCredentials::is_fresh) + { + Some(cached) => cached, + None => { + let assumed = + assume_bedrock_role(db, w_id, authed, role_arn, region, &session_name).await?; + BEDROCK_ASSUMED_ROLE_CACHE.insert(cache_key, assumed.clone()); + assumed + } + }; + + credentials.aws_access_key_id = Some(assumed.access_key_id); + credentials.aws_secret_access_key = Some(assumed.secret_access_key); + credentials.aws_session_token = Some(assumed.session_token); + + Ok(()) +} + +#[cfg(feature = "bedrock")] +async fn assume_bedrock_role( + db: &DB, + w_id: &str, + authed: &ApiAuthed, + role_arn: &str, + region: &str, + session_name: &str, +) -> Result { + #[cfg(all(feature = "enterprise", feature = "openidconnect", feature = "private"))] + { + use windmill_common::oidc_oss::{generate_id_token, WorkspaceClaim}; + + let id_token = generate_id_token( + Some(db), + WorkspaceClaim { workspace: w_id.to_string() }, + windmill_ai::ai_bedrock::AWS_OIDC_AUDIENCE, + format!("{}::{}", authed.email, w_id), + Some(authed.email.clone()), + None, + ) + .await?; + + windmill_ai::ai_bedrock::assume_role_with_oidc_token( + role_arn, + region, + id_token.to_string(), + session_name, + ) + .await + } + #[cfg(not(all(feature = "enterprise", feature = "openidconnect", feature = "private")))] + { + let _ = (db, w_id, authed, role_arn, region, session_name); + Err(Error::BadRequest( + "Assuming an AWS role through Windmill's OIDC provider requires an enterprise license" + .to_string(), + )) + } +} + #[derive(Clone, Debug)] pub struct ExpiringProviderCredentials { credentials: ProviderCredentials, @@ -976,6 +1107,7 @@ async fn global_proxy( aws_access_key_id: None, aws_secret_access_key: None, aws_session_token: None, + oidc_role_arn: None, platform: AIPlatform::Standard, custom_headers: HashMap::new(), }; @@ -1253,6 +1385,9 @@ async fn proxy( } }; + #[cfg(feature = "bedrock")] + resolve_bedrock_oidc_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. if free_lease.is_some() { @@ -1463,6 +1598,7 @@ mod tests { aws_access_key_id: None, aws_secret_access_key: None, aws_session_token: None, + oidc_role_arn: None, platform: AIPlatform::Standard, custom_headers: HashMap::new(), } @@ -1504,6 +1640,24 @@ mod tests { ("workspace-b".to_string(), AIProvider::OpenAI), ExpiringProviderCredentials::new(sample_provider_credentials(), None), ); + #[cfg(feature = "bedrock")] + let assumed_role_key = { + let key = ( + "workspace-a".to_string(), + "arn:aws:iam::123456789012:role/bedrock".to_string(), + "someone@windmill.dev".to_string(), + ); + BEDROCK_ASSUMED_ROLE_CACHE.insert( + key.clone(), + windmill_ai::ai_bedrock::AssumedRoleCredentials { + access_key_id: "AKIA".to_string(), + secret_access_key: "secret".to_string(), + session_token: "session".to_string(), + expires_at: std::time::SystemTime::now() + std::time::Duration::from_secs(3600), + }, + ); + key + }; invalidate_ai_request_cache_for_workspace("workspace-a"); @@ -1516,6 +1670,8 @@ mod tests { assert!(AI_REQUEST_CACHE .get(&("workspace-b".to_string(), AIProvider::OpenAI)) .is_some()); + #[cfg(feature = "bedrock")] + assert!(BEDROCK_ASSUMED_ROLE_CACHE.get(&assumed_role_key).is_none()); } #[test] diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 3c3a26875b..5f2ef1e6c0 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1419,6 +1419,11 @@ pub async fn run_agent( .map(|m| m.clamp(1, HARD_MAX_AGENT_ITERATIONS)) .unwrap_or(DEFAULT_MAX_AGENT_ITERATIONS); + // Held across iterations so an OIDC role is assumed once for the job and + // re-assumed only near expiry, rather than on every model turn. + #[cfg(feature = "bedrock")] + let mut bedrock_assumed_role: Option = None; + // Main agent loop for i in 0..max_iterations { // Check if parent was canceled — stop iterating but let current tool calls finish @@ -1438,6 +1443,27 @@ pub async fn run_agent( .region .as_deref() .unwrap_or(windmill_ai::ai_providers::USE_ENV_REGION); + // An OIDC role resolves to ordinary IAM keys, so from here the + // call is identical to the explicit-keys mode. + let assumed = windmill_ai::ai_bedrock::refresh_bedrock_oidc_credentials( + &credentials, + &mut bedrock_assumed_role, + client, + &job.id, + ) + .await?; + let (access_key_id, secret_access_key, session_token) = match assumed { + Some(assumed) => ( + Some(assumed.access_key_id.as_str()), + Some(assumed.secret_access_key.as_str()), + Some(assumed.session_token.as_str()), + ), + None => ( + credentials.aws_access_key_id.as_deref(), + credentials.aws_secret_access_key.as_deref(), + credentials.aws_session_token.as_deref(), + ), + }; // Use Bedrock SDK via dedicated query builder windmill_ai::providers::bedrock::BedrockQueryBuilder::default() .execute_request( @@ -1453,9 +1479,9 @@ pub async fn run_agent( client, &job.workspace_id, structured_output_tool_name.as_deref(), - credentials.aws_access_key_id.as_deref(), - credentials.aws_secret_access_key.as_deref(), - credentials.aws_session_token.as_deref(), + access_key_id, + secret_access_key, + session_token, ) .await? } diff --git a/frontend/src/lib/components/BedrockCredentialsCheck.svelte b/frontend/src/lib/components/BedrockCredentialsCheck.svelte index fa39309a50..726c6d8f0a 100644 --- a/frontend/src/lib/components/BedrockCredentialsCheck.svelte +++ b/frontend/src/lib/components/BedrockCredentialsCheck.svelte @@ -1,5 +1,6 @@