From ff2d7c985c568e371bf144e18b1da21c3e13e9f2 Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Tue, 10 Nov 2026 19:48:57 +0000 Subject: [PATCH] feat(messages): mint anthropic workload identity tokens on rust --- .../crates/inference-messages/src/handler.rs | 207 +++++++++++++++++- .../crates/inference-messages/src/lib.rs | 12 + .../src/workload_identity.rs | 11 + .../crates/llms/src/anthropic/common_utils.rs | 51 +++-- .../src/anthropic/messages/transformation.rs | 85 ++++--- .../src/base_llm/messages/transformation.rs | 30 +++ .../crates/python-bridge/src/errors.rs | 3 + .../python-bridge/src/routes/messages/host.rs | 68 +++++- .../python-bridge/src/routes/messages/mod.rs | 8 +- .../src/routes/messages/workload_identity.rs | 74 +++++++ litellm/messages/dispatch.py | 1 - litellm/rust_bridge/catalog.py | 1 - litellm/rust_bridge/host/routes/messages.py | 10 +- litellm/rust_bridge/messages/detectors.py | 18 -- .../sdk/messages/test_native_parity.py | 96 +++++++- .../messages/test_dispatch.py | 29 ++- .../rust_bridge/host/routes/test_messages.py | 36 +++ tests/unit/rust_bridge/test_detectors.py | 14 -- 18 files changed, 663 insertions(+), 91 deletions(-) create mode 100644 litellm-rust/crates/inference-messages/src/workload_identity.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/workload_identity.rs diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 46b16f1f79b..bb8d4f44367 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -18,7 +18,7 @@ use litellm_llms::base_llm::{ messages::{ context::MessagesTransformContext, streaming::{self, ByteStream, StreamDecoder}, - transformation::BaseMessagesConfig, + transformation::{BaseMessagesConfig, FederatedEnvironment}, }, }; use litellm_llms_types::formats::messages::{ @@ -29,7 +29,7 @@ use litellm_tracing::ByteChunk; use serde_json::{Map, Value}; use super::{ - Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, WorkloadIdentity, common_utils::{MessagesProvider, merge_messages_header_sources, string_headers}, constants::{MESSAGES_HTTP_ERROR_ATTEMPTS, MESSAGES_TIMEOUT_SECS}, route::Messages, @@ -105,6 +105,7 @@ pub(super) async fn execute( let caller_keyed = caller_keyed(route).is_some(); let request = provider_call( &route.auth, + route.workload_identity.as_deref(), provider, model, call, @@ -162,6 +163,7 @@ fn execution_facts( async fn provider_call( auth: &AuthServices, + workload_identity: Option<&dyn WorkloadIdentity>, provider: MessagesProvider, model: String, call: MessagesCall, @@ -203,14 +205,30 @@ async fn provider_call( extra_headers.unwrap_or_default(), scoped, ])))?; - let validated = config.validate_environment( + let validated = match config.validate_federated_environment( forwarded, api_key.as_deref(), &trimmed.model, &trimmed.params, &connection_params, &env_lookup, - )?; + )? { + FederatedEnvironment::Validated(validated) => validated, + FederatedEnvironment::Unauthenticated { + headers, + missing, + federate, + } => match workload_identity { + Some(identity) => match identity + .token(&connection_params, api_base.as_deref()) + .await? + { + Some(token) => federate(headers, token), + None => return Err(missing.into()), + }, + None => return Err(missing.into()), + }, + }; let tool_custom_field = config.tool_custom_field(&trimmed.model); let body = config.request_body( @@ -698,6 +716,14 @@ mod tests { async fn provider_call_with_secrets( call: MessagesCall, secrets: &(dyn Lookup + Send + Sync), + ) -> Result { + provider_call_federating(call, secrets, None).await + } + + async fn provider_call_federating( + call: MessagesCall, + secrets: &(dyn Lookup + Send + Sync), + workload_identity: Option<&dyn WorkloadIdentity>, ) -> Result { let provider = resolve_llm_provider( &call.body.model, @@ -711,6 +737,7 @@ mod tests { let context = CallContext::new(&interceptors, litellm_inference::CallOptions::default()); provider_call( &AuthServices::default(), + workload_identity, config, model, call, @@ -904,6 +931,178 @@ mod tests { assert!(error.is_missing_credential(), "{error:?}"); } + struct ScriptedIdentity { + minted: Result, ()>, + asked: std::sync::Mutex, Option)>>, + } + + impl ScriptedIdentity { + fn new(minted: Result, ()>) -> Self { + Self { + minted, + asked: std::sync::Mutex::new(Vec::new()), + } + } + + fn asked(&self) -> Vec<(Option, Option)> { + self.asked.lock().unwrap().clone() + } + } + + impl WorkloadIdentity for ScriptedIdentity { + fn token<'a>( + &'a self, + litellm_params: &'a litellm_router_types::LitellmParams, + api_base: Option<&'a str>, + ) -> futures_util::future::BoxFuture<'a, Result, litellm_auth::Error>> + { + self.asked.lock().unwrap().push(( + litellm_params + .extra + .get("anthropic_federation_rule_id") + .cloned(), + api_base.map(str::to_owned), + )); + let minted = self.minted; + Box::pin(async move { + minted + .map(|token| token.map(SecretValue::new)) + .map_err(|()| litellm_auth::Error::InvalidHeader) + }) + } + } + + fn federated_call( + shaping: MessagesShaping, + api_key: Option<&str>, + extra_headers: Option, + ) -> MessagesCall { + MessagesCall { + body: body( + json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), + ), + api_key: api_key.map(str::to_owned), + api_base: Some("https://anthropic.test".into()), + custom_llm_provider: Some("anthropic".into()), + litellm_params: crate::litellm_params( + serde_json::from_value( + json!({"model": "claude-test", "anthropic_federation_rule_id": "fdrl_1"}), + ) + .unwrap(), + ) + .unwrap(), + extra_headers: extra_headers.map(|headers| serde_json::from_value(headers).unwrap()), + provider_specific_header: None, + timeout: None, + shaping, + } + } + + fn sent_headers<'a>(call: &'a ProviderCall, names: &[&str]) -> Vec<(&'a str, &'a str)> { + call.wire + .headers + .iter() + .filter(|(name, _)| names.contains(&name.as_str())) + .map(|(name, value)| (name.as_str(), value.as_str())) + .collect() + } + + #[rstest] + #[case::alone(None, &[("anthropic-beta", "oauth-2025-04-20"), ("authorization", "Bearer sk-ant-oat01-minted")])] + #[case::beside_a_caller_beta( + Some(json!({"anthropic-beta": "files-api-2025-04-14"})), + &[("anthropic-beta", "files-api-2025-04-14,oauth-2025-04-20"), ("authorization", "Bearer sk-ant-oat01-minted")] + )] + #[tokio::test] + async fn a_minted_federation_token_is_sent_as_an_oauth_bearer( + shaping: MessagesShaping, + #[case] extra_headers: Option, + #[case] expected: &[(&str, &str)], + ) { + let identity = ScriptedIdentity::new(Ok(Some("sk-ant-oat01-minted"))); + let call = provider_call_federating( + federated_call(shaping, None, extra_headers), + &|_: &str| None, + Some(&identity), + ) + .await + .unwrap(); + let mut sent = sent_headers( + &call, + &[ + "anthropic-beta", + "authorization", + "x-api-key", + "anthropic-dangerous-direct-browser-access", + ], + ); + sent.sort_unstable(); + assert_eq!( + (sent.as_slice(), identity.asked()), + ( + expected, + vec![(Some(json!("fdrl_1")), Some("https://anthropic.test".into()))] + ) + ); + } + + #[rstest] + #[case::static_key(Some("sk-static"), None, &[("x-api-key", "sk-static")])] + #[case::caller_credential( + None, + Some(json!({"x-api-key": "sk-caller"})), + &[("x-api-key", "sk-caller")] + )] + #[tokio::test] + async fn a_configured_credential_is_never_federated( + shaping: MessagesShaping, + #[case] api_key: Option<&str>, + #[case] extra_headers: Option, + #[case] expected: &[(&str, &str)], + ) { + let identity = ScriptedIdentity::new(Ok(Some("sk-ant-oat01-minted"))); + let call = provider_call_federating( + federated_call(shaping, api_key, extra_headers), + &|_: &str| None, + Some(&identity), + ) + .await + .unwrap(); + assert_eq!( + ( + sent_headers(&call, &["authorization", "x-api-key"]), + identity.asked() + ), + (expected.to_vec(), Vec::new()) + ); + } + + #[rstest] + #[case::nothing_minted(Ok(None), true)] + #[case::minting_failed(Err(()), false)] + #[tokio::test] + async fn an_unminted_federation_token_fails_the_call( + shaping: MessagesShaping, + #[case] minted: Result, ()>, + #[case] missing_credential: bool, + ) { + let identity = ScriptedIdentity::new(minted); + let Err(error) = provider_call_federating( + federated_call(shaping, None, None), + &|_: &str| None, + Some(&identity), + ) + .await + else { + panic!("a call without a minted token fails"); + }; + assert_eq!( + error.is_missing_credential(), + missing_credential, + "{error:?}" + ); + } + #[rstest] #[case::nothing_configured(None, None, None)] #[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))] diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index ca9e0345ffe..5e8de57ed30 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -3,6 +3,7 @@ mod constants; mod handler; pub mod route; mod types; +mod workload_identity; use futures_util::FutureExt; use litellm_auth::AuthServices; @@ -20,6 +21,7 @@ pub use types::{ MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, messages_body, }; +pub use workload_identity::WorkloadIdentity; #[derive(Clone)] pub struct MessagesRoute { @@ -27,6 +29,7 @@ pub struct MessagesRoute { auth: Arc, secrets: Arc, cache: Option, + workload_identity: Option>, } impl MessagesRoute { @@ -40,6 +43,15 @@ impl MessagesRoute { auth, secrets, cache: None, + workload_identity: None, + } + } + + #[must_use] + pub fn with_workload_identity(self, workload_identity: Arc) -> Self { + Self { + workload_identity: Some(workload_identity), + ..self } } diff --git a/litellm-rust/crates/inference-messages/src/workload_identity.rs b/litellm-rust/crates/inference-messages/src/workload_identity.rs new file mode 100644 index 00000000000..56445de6eea --- /dev/null +++ b/litellm-rust/crates/inference-messages/src/workload_identity.rs @@ -0,0 +1,11 @@ +use futures_util::future::BoxFuture; +use litellm_router_types::LitellmParams; +use litellm_secret_value::SecretValue; + +pub trait WorkloadIdentity: Send + Sync { + fn token<'a>( + &'a self, + litellm_params: &'a LitellmParams, + api_base: Option<&'a str>, + ) -> BoxFuture<'a, Result, litellm_auth::Error>>; +} diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index f11416ea720..911ca0de881 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -22,7 +22,7 @@ use serde_json::Value; use crate::{ anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX, - base_llm::auth::{AuthScheme, Headers}, + base_llm::auth::{AuthScheme, Headers, ValidatedEnvironment}, }; pub const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; @@ -80,10 +80,21 @@ impl<'a> OauthToken<'a> { } pub fn into_auth(self) -> AuthScheme { - AuthScheme::Credential { - placement: CredentialPlacement::Bearer, - secret: SecretValue::new(self.0), - } + oauth_bearer(SecretValue::new(self.0)) + } +} + +fn oauth_bearer(secret: SecretValue) -> AuthScheme { + AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret, + } +} + +pub fn minted_oauth_environment(headers: Headers, token: SecretValue) -> ValidatedEnvironment { + ValidatedEnvironment { + headers: with_oauth_companions(headers, &[], BrowserAccess::Omitted), + auth: oauth_bearer(token), } } @@ -232,14 +243,24 @@ pub fn forwarded_oauth_bearer(headers: &[(String, String)]) -> Option Headers { - merge_beta_headers( +#[derive(Clone, Copy, PartialEq, Eq)] +enum BrowserAccess { + Granted, + Omitted, +} + +fn with_oauth_companions(headers: Headers, dropped: &[&str], browser: BrowserAccess) -> Headers { + let headers = merge_beta_headers( without_headers(headers, dropped), BetaSet::from_iter([AnthropicBeta::Oauth20250420]), - ) - .into_iter() - .chain([(DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string())]) - .collect() + ); + match browser { + BrowserAccess::Granted => headers + .into_iter() + .chain([(DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string())]) + .collect(), + BrowserAccess::Omitted => headers, + } } pub fn optionally_handle_anthropic_oauth(headers: Headers, api_key: Option<&str>) -> OauthHandling { @@ -247,13 +268,17 @@ pub fn optionally_handle_anthropic_oauth(headers: Headers, api_key: Option<&str> forwarded_oauth_bearer(&headers).map(|token| SecretValue::new(token.as_str())) { return OauthHandling::Bearer { - headers: with_oauth_companions(headers, &[API_KEY_HEADER, AUTHORIZATION]), + headers: with_oauth_companions( + headers, + &[API_KEY_HEADER, AUTHORIZATION], + BrowserAccess::Granted, + ), token, }; } if let Some(token) = api_key.and_then(OauthToken::parse) { return OauthHandling::Bearer { - headers: with_oauth_companions(headers, &[API_KEY_HEADER]), + headers: with_oauth_companions(headers, &[API_KEY_HEADER], BrowserAccess::Granted), token: SecretValue::new(token.as_str()), }; } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 2d39f444901..68390dbf5c9 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -26,13 +26,15 @@ use crate::{ anthropic_request_errors_on_thinking_binding_mismatch, complete_anthropic_url, existing_betas, get_auth_header, has_advisor_tool, has_anthropic_credential, is_anthropic_invalid_thinking_block_error, is_anthropic_thinking_binding_mismatch_error, - is_tool_search_used, merge_beta_headers, optionally_handle_anthropic_oauth, - requires_native_compaction_beta, strip_advisor_blocks, strip_encrypted_reasoning_blocks, - strip_thinking_blocks_from_anthropic_messages_request, + is_tool_search_used, merge_beta_headers, minted_oauth_environment, + optionally_handle_anthropic_oauth, requires_native_compaction_beta, strip_advisor_blocks, + strip_encrypted_reasoning_blocks, strip_thinking_blocks_from_anthropic_messages_request, }, base_llm::{ auth::AuthScheme, - messages::transformation::{BaseMessagesConfig, Headers, ValidatedEnvironment}, + messages::transformation::{ + BaseMessagesConfig, FederatedEnvironment, Headers, ValidatedEnvironment, + }, }, }; @@ -94,31 +96,22 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { _litellm_params: &LitellmParams, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - let headers = match optionally_handle_anthropic_oauth(headers, api_key) { - OauthHandling::Bearer { headers, token } => { - return Ok(ValidatedEnvironment { - headers, - auth: AuthScheme::Credential { - placement: CredentialPlacement::Bearer, - secret: token, - }, - }); - } - OauthHandling::Untouched(headers) => headers, - }; - if has_anthropic_credential(&headers) { - return Ok(ValidatedEnvironment { - headers, - auth: AuthScheme::Forwarded, - }); + match anthropic_environment(headers, api_key, env_lookup) { + FederatedEnvironment::Validated(validated) => Ok(validated), + FederatedEnvironment::Unauthenticated { missing, .. } => Err(missing), } - let auth = get_auth_header(api_key, env_lookup).ok_or(Error::Auth( - litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: ANTHROPIC_API_KEY_ENV, - }, - ))?; - Ok(ValidatedEnvironment { headers, auth }) + } + + fn validate_federated_environment( + &self, + headers: Headers, + api_key: Option<&str>, + _model: &str, + _optional_params: &MessagesOptionalParams, + _litellm_params: &LitellmParams, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(anthropic_environment(headers, api_key, env_lookup)) } fn default_headers(&self) -> &'static [(&'static str, &'static str)] { @@ -233,6 +226,42 @@ pub fn retry_anthropic_messages_after_http_error( ) } +fn anthropic_environment( + headers: Headers, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> FederatedEnvironment { + let headers = match optionally_handle_anthropic_oauth(headers, api_key) { + OauthHandling::Bearer { headers, token } => { + return FederatedEnvironment::Validated(ValidatedEnvironment { + headers, + auth: AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: token, + }, + }); + } + OauthHandling::Untouched(headers) => headers, + }; + if has_anthropic_credential(&headers) { + return FederatedEnvironment::Validated(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + match get_auth_header(api_key, env_lookup) { + Some(auth) => FederatedEnvironment::Validated(ValidatedEnvironment { headers, auth }), + None => FederatedEnvironment::Unauthenticated { + headers, + missing: Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: ANTHROPIC_API_KEY_ENV, + }), + federate: minted_oauth_environment, + }, + } +} + pub(crate) fn transform_messages_request( request: MessagesRequest, context: &MessagesTransformContext, diff --git a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs index 72475e9c316..13c021c2aec 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -3,6 +3,7 @@ use litellm_llms_types::formats::messages::{ MessagesOptionalParams, MessagesRequest, MessagesResponse, }; use litellm_router_types::LitellmParams; +use litellm_secret_value::SecretValue; use serde_json::Value; use super::{ @@ -13,6 +14,15 @@ use super::{ pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; use crate::{Error, base_llm::messages::streaming::StreamDecoder}; +pub enum FederatedEnvironment { + Validated(ValidatedEnvironment), + Unauthenticated { + headers: Headers, + missing: Error, + federate: fn(Headers, SecretValue) -> ValidatedEnvironment, + }, +} + pub trait BaseMessagesConfig: Sync { fn shape_request( &self, @@ -70,6 +80,26 @@ pub trait BaseMessagesConfig: Sync { env_lookup: &dyn Fn(&str) -> Option, ) -> Result; + fn validate_federated_environment( + &self, + headers: Headers, + api_key: Option<&str>, + model: &str, + optional_params: &MessagesOptionalParams, + litellm_params: &LitellmParams, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + self.validate_environment( + headers, + api_key, + model, + optional_params, + litellm_params, + env_lookup, + ) + .map(FederatedEnvironment::Validated) + } + /// `None` relays the upstream bytes untouched, which is right for every host that already /// speaks Anthropic SSE. A host on another wire returns the decoder that lifts its frames /// into Anthropic stream events, and the route re-encodes those as Anthropic SSE. diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index db5696d52ce..0ba6697a877 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -90,6 +90,9 @@ pub(crate) fn route_failure( if let Some(raised) = crate::secrets::python_error(py, &error) { return Ok(raised); } + if let Some(raised) = crate::routes::messages::workload_identity::python_error(py, &error) { + return Ok(raised); + } if let Some((status, body)) = error.upstream_response() { let headers = error .upstream_headers() diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 6019bae696d..2a67e7e70e4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -14,6 +14,7 @@ use pyo3::{ exceptions::PyException, gc::{PyTraverseError, PyVisit}, prelude::*, + sync::PyOnceLock, types::{PyBytes, PyDict}, }; use serde_json::{Map, Value}; @@ -40,6 +41,19 @@ fn merge_headers( (!merged.is_empty()).then_some(merged) } +const FEDERATING_PROVIDER: &str = "anthropic"; + +fn workload_identity_fields(py: Python<'_>) -> PyResult<&'static [String]> { + static FIELDS: PyOnceLock> = PyOnceLock::new(); + FIELDS + .get_or_try_init(py, || { + py.import(ROUTE_HOST_MODULE)? + .getattr("WORKLOAD_IDENTITY_FIELDS")? + .extract() + }) + .map(Vec::as_slice) +} + /// Reads `litellm.`, the module global a param spec names, the way /// `VertexBase.get_vertex_ai_project` reads `litellm.vertex_project`. fn module_global<'py>(py: Python<'py>, name: &str) -> PyResult>> { @@ -55,11 +69,17 @@ fn module_global<'py>(py: Python<'py>, name: &str) -> PyResult( +fn project_litellm_params<'a, 'py>( + names: impl IntoIterator, argument: impl Fn(&str) -> PyResult>>, global: impl Fn(&str) -> PyResult>>, ) -> PyResult> { - let fields = project_optional_fields(LitellmParams::fields(), &argument)?; + let fields = project_optional_fields( + LitellmParams::fields() + .map(|name| -> &'a str { name }) + .chain(names), + &argument, + )?; let globals = LitellmParams::specs() .filter(|spec| !spec.wire.iter().any(|name| fields.contains_key(*name))) .filter_map(|spec| spec.module_global); @@ -175,7 +195,15 @@ impl MessagesPythonHost { None => self.merged_headers(py, arguments)?, }; let provider_specific_header = self.provider_specific_header(py, arguments)?; - let litellm_params = project_litellm_params(argument, |name| module_global(py, name))?; + let federation_fields: &[String] = match custom_llm_provider.as_deref() { + Some(FEDERATING_PROVIDER) => workload_identity_fields(py)?, + _ => &[], + }; + let litellm_params = project_litellm_params( + federation_fields.iter().map(String::as_str), + argument, + |name| module_global(py, name), + )?; Ok(messages_body(body).and_then(|body| { Ok(MessagesCall { body, @@ -482,7 +510,40 @@ mod tests { assert_projection(kwargs, globals, expected); } + #[rstest] + #[case::named_fields_are_projected( + json!({"model": "m", "anthropic_federation_rule_id": "fdrl_1", "anthropic_issuer_ttl_seconds": 60}), + Ok(json!({"model": "m", "anthropic_federation_rule_id": "fdrl_1", "anthropic_issuer_ttl_seconds": 60})), + )] + #[case::unnamed_fields_are_not( + json!({"model": "m", "anthropic_unrelated": "x"}), + Ok(json!({"model": "m"})), + )] + fn workload_identity_fields_are_projected_by_the_names_the_host_lists( + #[case] kwargs: Value, + #[case] expected: Result, + ) { + assert_projection_of( + &[ + "anthropic_federation_rule_id", + "anthropic_issuer_ttl_seconds", + ], + kwargs, + json!({}), + expected, + ); + } + fn assert_projection(kwargs: Value, globals: Value, expected: Result) { + assert_projection_of(&[], kwargs, globals, expected); + } + + fn assert_projection_of( + names: &[&str], + kwargs: Value, + globals: Value, + expected: Result, + ) { Python::initialize(); Python::attach(|py| { let dict = |value: &Value| { @@ -496,6 +557,7 @@ mod tests { let globals = dict(&globals); let bound = PyDict::new(py); let projected = project_litellm_params( + names.iter().copied(), |name| present(&kwargs, &bound, name), |name| present(&globals, &bound, name), ) diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index 7c67db064e3..9b61beb49ed 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,5 +1,6 @@ mod host; mod modify_params; +pub(crate) mod workload_identity; use host::MessagesPythonHost; use litellm_callbacks_legacy_python::LoggingOperation; @@ -21,7 +22,12 @@ fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyR .map_err(crate::http::client_error)?, crate::http::resources().auth.clone(), crate::secrets::source(py)?, - ); + ) + .with_workload_identity(std::sync::Arc::new( + workload_identity::PythonWorkloadIdentity::new( + litellm_host_python::PythonContext::capture(py)?, + ), + )); Ok(litellm_host::call::hosted_call( request, None, diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/workload_identity.rs b/litellm-rust/crates/python-bridge/src/routes/messages/workload_identity.rs new file mode 100644 index 00000000000..b8824d96584 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/workload_identity.rs @@ -0,0 +1,74 @@ +use std::fmt; + +use futures_util::future::BoxFuture; +use litellm_host_python::{PythonContext, attach_blocking, to_py}; +use litellm_inference_messages::WorkloadIdentity; +use litellm_router_types::{ExtraParams, LitellmParams}; +use litellm_secret_value::SecretValue; +use pyo3::{exceptions::PyBaseException, prelude::*}; + +const MODULE: &str = "litellm.rust_bridge.host.routes.messages"; + +pub(super) struct PythonWorkloadIdentity { + context: PythonContext, +} + +impl PythonWorkloadIdentity { + pub(super) fn new(context: PythonContext) -> Self { + Self { context } + } +} + +impl WorkloadIdentity for PythonWorkloadIdentity { + fn token<'a>( + &'a self, + litellm_params: &'a LitellmParams, + api_base: Option<&'a str>, + ) -> BoxFuture<'a, Result, litellm_auth::Error>> { + let params = litellm_params.extra.clone(); + let api_base = api_base.map(str::to_owned); + let context = self.context.clone(); + Box::pin(async move { + attach_blocking(context, move |py| { + mint(py, ¶ms, api_base).map_err(|error| raised(py, error)) + }) + .await + .map_err(|error| Python::attach(|py| raised(py, error))) + .and_then(|minted| minted) + .map(|token| token.map(SecretValue::new)) + }) + } +} + +fn mint( + py: Python<'_>, + params: &ExtraParams, + api_base: Option, +) -> PyResult> { + py.import(MODULE)? + .getattr("workload_identity_token")? + .call1((to_py(py, &**params)?, api_base))? + .extract() +} + +#[derive(thiserror::Error)] +#[error("Python workload identity federation failed")] +struct PythonWorkloadIdentityError(Py); + +impl fmt::Debug for PythonWorkloadIdentityError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("PythonWorkloadIdentityError") + } +} + +fn raised(py: Python<'_>, error: PyErr) -> litellm_auth::Error { + litellm_auth::Error::provider(PythonWorkloadIdentityError(error.into_value(py))) +} + +pub(crate) fn python_error( + py: Python<'_>, + error: &(dyn std::error::Error + 'static), +) -> Option { + litellm_error::find::(error) + .map(|error| PyErr::from_value(error.0.clone_ref(py).into_bound(py).into_any())) +} diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index c1bfbdf1c85..c41e0d9ed17 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -139,7 +139,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType( { **{feature: detect for feature, detect in detectors.COMMON.items() if feature not in _NATIVE_COMMON_FEATURES}, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, - Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity, } ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index da011a69fcd..e72c9d05a79 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -169,7 +169,6 @@ MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VAR Feature.PRE_REQUEST_HOOK, Feature.WEBSEARCH_INTERCEPTION, Feature.MESSAGES_INTERCEPTOR, - Feature.WORKLOAD_IDENTITY, Feature.STRICT_TOOLS_RETRY, } diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index ae4acb9ceb8..e6ddf7d2742 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -43,6 +43,7 @@ from litellm.llms.anthropic.messages.native.handler import ( ) from litellm.llms.anthropic.messages.native.utils import prepare_native_messages from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled +from litellm.llms.anthropic.wif import get_anthropic_wif_token from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.llms.custom_httpx.llm_http_handler import ( @@ -53,7 +54,7 @@ from litellm.rust_bridge.host import cache, callbacks_legacy_python, http_transp from litellm.rust_bridge.public_call import NativeCall, known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import CallTypes, ProviderSpecificHeader +from litellm.types.utils import CallTypes, ProviderSpecificHeader, anthropic_wif_litellm_params from litellm.utils import clamp_max_tokens _METADATA: Final = TypeAdapter(dict[str, object]) @@ -85,6 +86,8 @@ _BODY_FIELDS: Final = frozenset( } ) +WORKLOAD_IDENTITY_FIELDS: Final = anthropic_wif_litellm_params + @dataclass(frozen=True, slots=True) class MessagesSettings: @@ -635,3 +638,8 @@ def inject_stream_cost(frames: bytes, model: str, request: Mapping[str, object]) as_str_mapping(logging_obj.model_call_details.get("litellm_params")) if logging_obj else None ), ) + + +def workload_identity_token(params: Mapping[str, object], api_base: str | None) -> str | None: + litellm_params: Final = GenericLiteLLMParams(**params) + return get_anthropic_wif_token(dict(litellm_params), api_base, "") diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 0a14ea08af0..78d6d601fb0 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -14,24 +14,6 @@ def streamed_mock_response(request: NativeCall) -> bool: return present(request, "mock_response") and request.resolved.get("stream") is True -def anthropic_workload_identity(request: NativeCall) -> bool: - """Python mints an Anthropic federation token when no static credential is configured.""" - from litellm.llms.anthropic.common_utils import AnthropicModelInfo - from litellm.secret_managers.main import get_secret_str - - def configured(param: str, env: str) -> bool: - return present(request, param) or get_secret_str(env) is not None - - return ( - AnthropicModelInfo.get_api_key(optional_str(request.resolved.get("api_key"))) is None - and AnthropicModelInfo.get_auth_token() is None - and configured("anthropic_federation_rule_id", "ANTHROPIC_FEDERATION_RULE_ID") - and configured("anthropic_organization_id", "ANTHROPIC_ORGANIZATION_ID") - ) - - - - def strict_tools_retry(request: NativeCall) -> bool: tools: Final = optional_sequence(request.resolved.get("tools")) or () return drops_unsupported_params(request.resolved.get("drop_params")) and any( diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py index f8fbdc210ff..6e496712e0e 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py @@ -39,6 +39,7 @@ from litellm.caching.caching_handler import ( _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # a primed entry must land before the measured call reads it ) from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.llms.base_llm.auth.token_exchange import default_token_exchange_engine from litellm.messages import dispatch from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.rust_bridge import catalog @@ -1090,6 +1091,76 @@ CACHE_CASES: Final = ( Case("cache_sync", sync=True, upstream=(_ok(), _ok()), cache=Cache), ) +FEDERATION_ENV: Final[Mapping[str, str]] = MappingProxyType( + {"ANTHROPIC_IDENTITY_TOKEN": "eyJ.parity-identity.jwt", "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS": "127.0.0.1"} +) +FEDERATION_IDS: Final[Mapping[str, str]] = MappingProxyType( + {"anthropic_federation_rule_id": "fdrl_parity", "anthropic_organization_id": "org-parity"} +) +MINTED: Final = ResponseSpec( + body={"access_token": "sk-ant-oat01-parity-minted", "token_type": "Bearer", "expires_in": 3600} +) + +WORKLOAD_IDENTITY_CASES: Final = ( + Case( + "workload_identity_from_params", + _kw(api_key=None, **FEDERATION_IDS), + upstream=(MINTED, _ok()), + env=FEDERATION_ENV, + ), + Case( + "workload_identity_from_environment", + _kw(api_key=None), + upstream=(MINTED, _ok()), + env=MappingProxyType( + {**FEDERATION_ENV, "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_env", "ANTHROPIC_ORGANIZATION_ID": "org-env"} + ), + ), + Case( + "workload_identity_with_workspace_and_caller_beta", + _kw( + api_key=None, + anthropic_federation_workspace_id="wrkspc_parity", + anthropic_service_account_id="svac_parity", + extra_headers={"anthropic-beta": "files-api-2025-04-14", "x-trace-id": "trace-1"}, + **FEDERATION_IDS, + ), + upstream=(MINTED, _ok()), + env=FEDERATION_ENV, + ), + Case( + "workload_identity_streamed", + _kw(api_key=None, stream=True, **FEDERATION_IDS), + upstream=(MINTED, _stream(TEXT_STREAM)), + env=FEDERATION_ENV, + ), + Case( + "workload_identity_exchange_denied", + _kw(api_key=None, **FEDERATION_IDS), + upstream=(_failure(401, "authentication_error", "invalid assertion"),), + env=FEDERATION_ENV, + ), + Case( + "workload_identity_untrusted_exchange_host", + _kw(api_key=None, **FEDERATION_IDS), + upstream=(), + env=MappingProxyType({"ANTHROPIC_IDENTITY_TOKEN": "eyJ.parity-identity.jwt"}), + ), + Case( + "workload_identity_requested_without_ids", + _kw(api_key=None, anthropic_identity_token="oidc/env/ANTHROPIC_IDENTITY_TOKEN"), + upstream=(), + env=FEDERATION_ENV, + ), + Case( + "workload_identity_disabled", + _kw(api_key=None, anthropic_disable_workload_identity_federation=True, **FEDERATION_IDS), + upstream=(), + env=FEDERATION_ENV, + ), + Case("workload_identity_shadowed_by_a_static_key", _kw(**FEDERATION_IDS), env=FEDERATION_ENV), +) + CASES: Final = ( *REQUEST_CASES, *PROMPT_CACHING_CASES, @@ -1102,6 +1173,7 @@ CASES: Final = ( *THINKING_RETRY_CASES, *FALLBACK_CASES, *CACHE_CASES, + *WORKLOAD_IDENTITY_CASES, ) @@ -1154,6 +1226,12 @@ def _environment(env: Mapping[str, str]) -> Generator[None]: *env, "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_FEDERATION_RULE_ID", + "ANTHROPIC_ORGANIZATION_ID", + "ANTHROPIC_IDENTITY_TOKEN", + "ANTHROPIC_IDENTITY_TOKEN_FILE", + "ANTHROPIC_IDENTITY_SOURCE", + "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", "LITELLM_REASONING_AUTO_SUMMARY", "LITELLM_RUST", ) @@ -1186,6 +1264,16 @@ def _cache(cache: Callable[[], Cache] | None) -> Generator[None]: yield +@contextmanager +def _fresh_token_exchange() -> Generator[None]: + # The federation token cache and its on-disk store would hand the second leg the first leg's token. + with ( + rebound(default_token_exchange_engine, "_entries", {}), + rebound(default_token_exchange_engine, "_shared_store", None), + ): + yield + + def _kwargs(case: Case, server: RecordingServer) -> dict[str, object]: # Deep: Python writes into the caller's own `litellm_metadata`, which must not leak into the # other leg. @@ -1363,7 +1451,13 @@ async def _prime(case: Case, server: RecordingServer, recorder: RecordingLogger, async def _leg(case: Case, server: RecordingServer, probe: FallbackProbe | None) -> Leg: - with isolated_callback_registries(), _environment(case.env), _settings(case.settings), _cache(case.cache): + with ( + isolated_callback_registries(), + _environment(case.env), + _settings(case.settings), + _cache(case.cache), + _fresh_token_exchange(), + ): for spec in case.upstream: server.enqueue(replace(spec)) request_count: Final = len(server.requests) diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index c44b810b255..933b92095eb 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -478,12 +478,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": (Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}), (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), - ( - Feature.WORKLOAD_IDENTITY, - "anthropic/claude-sonnet-4-5", - {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"}, - {}, - ), ), ) def test_each_detected_feature_reaches_its_rule( @@ -847,3 +841,26 @@ async def test_interceptor_feature_follows_the_resolved_provider() -> None: rules=rules, ) assert result is native_result + + +def test_a_federated_request_stays_native(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"): + monkeypatch.delenv(name, raising=False) + rules: Final[Rules] = ( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.WORKLOAD_IDENTITY})), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + native_result: Final = response() + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("Rust mints the federation token itself") + + served: Final = _DISPATCH.run( + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"}, + python=python, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, call_args, native_kwargs: hook(request), + rules=rules, + ) + assert served is native_result diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index cbaa13d0cf8..ea739291ff4 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -651,3 +651,39 @@ def test_modify_max_tokens_clamps_the_keyword_view_as_the_client_wrapper_does( messages_host.modify_max_tokens(model, args, view) assert view["max_tokens"] == expected + + +def test_workload_identity_token_mints_with_the_params_python_validates(monkeypatch: pytest.MonkeyPatch) -> None: + minted: Final[list[tuple[dict[str, object], str | None, str]]] = [] + + def mint(litellm_params: dict[str, object], api_base: str | None, model: str) -> str: + federation = {name: value for name, value in litellm_params.items() if name.startswith("anthropic_")} + minted.append(({name: value for name, value in federation.items() if value is not None}, api_base, model)) + return "sk-ant-oat01-minted" + + monkeypatch.setattr(messages_host, "get_anthropic_wif_token", mint) + + token: Final = messages_host.workload_identity_token( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_issuer_ttl_seconds": "60"}, + "https://gateway.test", + ) + + assert (token, minted) == ( + "sk-ant-oat01-minted", + [({"anthropic_federation_rule_id": "fdrl_1", "anthropic_issuer_ttl_seconds": 60}, "https://gateway.test", "")], + ) + + +def test_workload_identity_token_is_none_without_a_federation_rule(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("ANTHROPIC_FEDERATION_RULE_ID", "ANTHROPIC_ORGANIZATION_ID", "ANTHROPIC_IDENTITY_TOKEN"): + monkeypatch.delenv(name, raising=False) + + assert messages_host.workload_identity_token({}, None) is None + + +def test_workload_identity_fields_are_the_federation_litellm_params() -> None: + assert set(messages_host.WORKLOAD_IDENTITY_FIELDS) >= { + "anthropic_federation_rule_id", + "anthropic_organization_id", + "anthropic_disable_workload_identity_federation", + } diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 1b7bc1c848e..80d877930aa 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -369,20 +369,6 @@ def test_openai_workload_identity_applies_only_without_a_static_key(monkeypatch: assert not responses_detectors.openai_workload_identity(call()) -def test_anthropic_workload_identity_applies_only_without_a_static_key(monkeypatch: pytest.MonkeyPatch) -> None: - for name in ( - "ANTHROPIC_API_KEY", - "ANTHROPIC_AUTH_TOKEN", - "ANTHROPIC_FEDERATION_RULE_ID", - "ANTHROPIC_ORGANIZATION_ID", - ): - monkeypatch.delenv(name, raising=False) - federation: Final = {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"} - assert messages_detectors.anthropic_workload_identity(call(**federation)) - assert not messages_detectors.anthropic_workload_identity(call(**federation, api_key="sk-ant-static")) - assert not messages_detectors.anthropic_workload_identity(call(anthropic_federation_rule_id="rule-1")) - - class _Logger(CustomLogger): pass