diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index c7a7d93bf98..b952f3c42f1 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_TIMEOUT_SECS, route::Messages, @@ -62,6 +62,7 @@ pub(super) async fn execute( ) -> Result { let request = provider_call( &route.auth, + route.workload_identity.as_deref(), provider, model, call, @@ -103,6 +104,7 @@ pub(super) async fn execute( async fn provider_call( auth: &AuthServices, + workload_identity: Option<&dyn WorkloadIdentity>, provider: MessagesProvider, model: String, call: MessagesCall, @@ -144,14 +146,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 body = config.request_body( config.transform_anthropic_messages_request(trimmed, &transform_context)?, @@ -529,6 +547,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, @@ -542,6 +568,7 @@ mod tests { let context = CallContext::new(&interceptors, litellm_inference::CallOptions::default()); provider_call( &AuthServices::default(), + workload_identity, config, model, call, @@ -735,6 +762,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 25ebeaacb49..f5238b0a2a6 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 0854820358b..7ab891b5517 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 18789e5b2be..3ef4e779fed 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -18,12 +18,14 @@ use crate::{ ANTHROPIC_API_BASE_ENV, ANTHROPIC_API_KEY_ENV, ANTHROPIC_AUTH_TOKEN_ENV, ANTHROPIC_BASE_URL_ENV, OauthHandling, complete_anthropic_url, get_auth_header, has_advisor_tool, has_anthropic_credential, is_tool_search_used, merge_beta_headers, - optionally_handle_anthropic_oauth, requires_native_compaction_beta, strip_advisor_blocks, - strip_encrypted_reasoning_blocks, + minted_oauth_environment, optionally_handle_anthropic_oauth, + requires_native_compaction_beta, strip_advisor_blocks, strip_encrypted_reasoning_blocks, }, base_llm::{ auth::AuthScheme, - messages::transformation::{BaseMessagesConfig, Headers, ValidatedEnvironment}, + messages::transformation::{ + BaseMessagesConfig, FederatedEnvironment, Headers, ValidatedEnvironment, + }, }, }; @@ -79,31 +81,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)] { @@ -118,6 +111,42 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { } } +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 744ce89afdb..4bb39edfae0 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -2,6 +2,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::context::MessagesTransformContext; @@ -9,6 +10,15 @@ use super::context::MessagesTransformContext; 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, @@ -62,6 +72,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 2a5f79ca393..71dbc0e94e6 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); @@ -120,7 +140,15 @@ impl MessagesPythonHost { let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; let extra_headers = 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, @@ -421,7 +449,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| { @@ -435,6 +496,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 a9ae01bd7ee..121cf6658fc 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,4 +1,5 @@ mod host; +pub(crate) mod workload_identity; use host::MessagesPythonHost; use litellm_callbacks_legacy_python::LoggingOperation; @@ -18,7 +19,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 faa604a25ce..1a68a365e5b 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -140,7 +140,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType( Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, Feature.RESPONSE_CACHE: detectors.response_cache, Feature.MODIFY_PARAMS: detectors.modify_params, - Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity, Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection, Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection, Feature.THINKING_SIGNATURE_RETRY: route_detectors.thinking_signature_retry, diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index c493617003d..acc82f30a3d 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -181,7 +181,6 @@ MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VAR Feature.MESSAGES_INTERCEPTOR, Feature.RESPONSE_CACHE, Feature.MODIFY_PARAMS, - Feature.WORKLOAD_IDENTITY, Feature.CACHE_CONTROL_INJECTION, Feature.STREAM_COST_INJECTION, Feature.THINKING_SIGNATURE_RETRY, diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index c6568264edb..4ff1ec0c1dd 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -28,12 +28,15 @@ from litellm.llms.anthropic.messages.native.handler import ( _with_deployment_output_effort as _deployment_output_effort, # pyright: ignore[reportPrivateUsage] # the handler's own effort rounding ) 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.rust_bridge import failures from litellm.rust_bridge.host import callbacks_legacy_python, http_transport from litellm.rust_bridge.public_call import 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 anthropic_wif_litellm_params _METADATA: Final = TypeAdapter(dict[str, object]) _BODY_FIELDS: Final = frozenset( @@ -64,6 +67,8 @@ _BODY_FIELDS: Final = frozenset( } ) +WORKLOAD_IDENTITY_FIELDS: Final = anthropic_wif_litellm_params + @dataclass(frozen=True, slots=True) class MessagesSettings: @@ -394,3 +399,8 @@ def settings(kwargs: Mapping[str, object]) -> dict[str, object]: request_timeout=get_configured_request_timeout(), ) ) + + +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 99555bb4852..f3ea31f2c87 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -12,7 +12,7 @@ from litellm.litellm_core_utils.core_helpers import ( stream_cost_injection_requested, ) from litellm.llms.opencode.common_utils import is_opencode_api_base -from litellm.rust_bridge.detectors import api_base, overridden, present, streaming +from litellm.rust_bridge.detectors import api_base, overridden, streaming from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str @@ -26,22 +26,6 @@ def foundry_beta_row(request: NativeCall) -> bool: return is_foundry_api_base(AnthropicModelInfo.get_api_base(api_base(request))) -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") - ) - - _THINKING_BLOCKS: Final = frozenset({"thinking", "redacted_thinking"}) 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 41008f0d57f..53c87169de9 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 @@ -35,6 +35,7 @@ from pydantic import JsonValue, TypeAdapter import litellm 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.rust_bridge import catalog from litellm.rust_bridge.catalog import Route, RouteRule @@ -769,7 +770,80 @@ FALLBACK_CASES: Final = ( Case("stream_timeout_kwarg", _kw(stream=True, stream_timeout=30), upstream=(_stream(TEXT_STREAM),)), ) -CASES: Final = (*REQUEST_CASES, *RESPONSE_CASES, *FAILURE_CASES, *STREAM_CASES, *SYNC_CASES, *FALLBACK_CASES) +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, + *RESPONSE_CASES, + *FAILURE_CASES, + *STREAM_CASES, + *SYNC_CASES, + *FALLBACK_CASES, + *WORKLOAD_IDENTITY_CASES, +) @dataclass(frozen=True, slots=True) @@ -820,6 +894,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", ) @@ -843,6 +923,13 @@ def _settings(settings: Mapping[str, object]) -> Generator[None]: setattr(litellm, name, value) +@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. @@ -1001,7 +1088,12 @@ def _logged(recorder: RecordingLogger, events: tuple[str, ...]) -> tuple[Mapping async def _leg(case: Case, server: RecordingServer, probe: FallbackProbe | None) -> Leg: - with isolated_callback_registries(), _environment(case.env), _settings(case.settings): + with ( + isolated_callback_registries(), + _environment(case.env), + _settings(case.settings), + _fresh_token_exchange(), + ): recorder: Final = RecordingLogger() litellm.callbacks.append(recorder) for spec in case.upstream: diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index b580eca2e42..9ab062c925e 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -479,12 +479,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": (Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}), (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), (Feature.MODIFY_PARAMS, "anthropic/claude-sonnet-4-5", {}, {"modify_params": True}), - ( - Feature.WORKLOAD_IDENTITY, - "anthropic/claude-sonnet-4-5", - {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"}, - {}, - ), (Feature.CACHE_CONTROL_INJECTION, "anthropic/claude-sonnet-4-5", {"enable_prompt_caching": True}, {}), ( Feature.THINKING_SIGNATURE_RETRY, @@ -709,3 +703,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 aa7dcc74841..b48c267f7b7 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -445,3 +445,39 @@ def test_a_failure_before_any_response_maps_like_the_python_transport_error( f"litellm.InternalServerError: AnthropicException - {text}. Handle with `litellm.InternalServerError`." ) assert "This looks like a bug in LiteLLM" not in str(mapped) + + +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 fccc01dd8e6..afd8d9a8439 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -413,20 +413,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