scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

feat(messages): mint anthropic workload identity tokens on rust

Browse source
This commit is contained in:
narmstrong-berry 2026-11-10 19:48:57 +00:00
parent f0d234c4c6
commit 01ab88710e
No known key found for this signature in database
18 changed files with 663 additions and 89 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

207
litellm-rust/crates/inference-messages/src/handler.rs
View file

@ -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<MessagesCallResponse, Error> {
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<ProviderCall, Error> {
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<ProviderCall, Error> {
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<Option<&'static str>, ()>,
asked: std::sync::Mutex<Vec<(Option<Value>, Option<String>)>>,
}
impl ScriptedIdentity {
fn new(minted: Result<Option<&'static str>, ()>) -> Self {
Self {
minted,
asked: std::sync::Mutex::new(Vec::new()),
}
}
fn asked(&self) -> Vec<(Option<Value>, Option<String>)> {
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<Option<SecretValue>, 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<Value>,
) -> 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<Value>,
#[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<Value>,
#[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<Option<&'static str>, ()>,
#[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)))]

12
litellm-rust/crates/inference-messages/src/lib.rs
View file

@ -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<AuthServices>,
secrets: Arc<dyn SecretSource>,
cache: Option<litellm_cache_response::ScopedCache>,
workload_identity: Option<Arc<dyn WorkloadIdentity>>,
}
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<dyn WorkloadIdentity>) -> Self {
Self {
workload_identity: Some(workload_identity),
..self
}
}

11
litellm-rust/crates/inference-messages/src/workload_identity.rs Normal file
View file

@ -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<Option<SecretValue>, litellm_auth::Error>>;
}

51
litellm-rust/crates/llms/src/anthropic/common_utils.rs
View file

@ -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<OauthToken
.and_then(OauthToken::parse)
}
fn with_oauth_companions(headers: Headers, dropped: &[&str]) -> 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()),
};
}

83
litellm-rust/crates/llms/src/anthropic/messages/transformation.rs
View file

@ -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<String>,
) -> Result<ValidatedEnvironment, Error> {
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<String>,
) -> Result<FederatedEnvironment, Error> {
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<String>,
) -> 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,

30
litellm-rust/crates/llms/src/base_llm/messages/transformation.rs
View file

@ -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<String>,
) -> Result<ValidatedEnvironment, Error>;
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<String>,
) -> Result<FederatedEnvironment, Error> {
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.

3
litellm-rust/crates/python-bridge/src/errors.rs
View file

@ -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()

68
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
View file

@ -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<Vec<String>> = PyOnceLock::new();
FIELDS
.get_or_try_init(py, || {
py.import(ROUTE_HOST_MODULE)?
.getattr("WORKLOAD_IDENTITY_FIELDS")?
.extract()
})
.map(Vec::as_slice)
}
/// Reads `litellm.<name>`, 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<Option<Bound<'py, PyAny>>> {
@ -55,11 +69,17 @@ fn module_global<'py>(py: Python<'py>, name: &str) -> PyResult<Option<Bound<'py,
/// The caller's litellm params, read from the kwargs by the names the typed params declare,
/// so a key the configs do not read is never converted; a spec whose spellings the call all
/// leaves out falls back to the module global it names, which stays the host's concern.
fn project_litellm_params<'py>(
fn project_litellm_params<'a, 'py>(
names: impl IntoIterator<Item = &'a str>,
argument: impl Fn(&str) -> PyResult<Option<Bound<'py, PyAny>>>,
global: impl Fn(&str) -> PyResult<Option<Bound<'py, PyAny>>>,
) -> PyResult<Result<LitellmParams, Error>> {
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<Value, ()>,
) {
assert_projection_of(
&[
"anthropic_federation_rule_id",
"anthropic_issuer_ttl_seconds",
],
kwargs,
json!({}),
expected,
);
}
fn assert_projection(kwargs: Value, globals: Value, expected: Result<Value, ()>) {
assert_projection_of(&[], kwargs, globals, expected);
}
fn assert_projection_of(
names: &[&str],
kwargs: Value,
globals: Value,
expected: Result<Value, ()>,
) {
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),
)

8
litellm-rust/crates/python-bridge/src/routes/messages/mod.rs
View file

@ -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,

74
litellm-rust/crates/python-bridge/src/routes/messages/workload_identity.rs Normal file
View file

@ -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<Option<SecretValue>, 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, &params, 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<String>,
) -> PyResult<Option<String>> {
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<PyBaseException>);
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<PyErr> {
litellm_error::find::<PythonWorkloadIdentityError>(error)
.map(|error| PyErr::from_value(error.0.clone_ref(py).into_bound(py).into_any()))
}

1
litellm/messages/dispatch.py
View file

@ -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,

1
litellm/rust_bridge/catalog.py
View file

@ -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,

10
litellm/rust_bridge/host/routes/messages.py
View file

@ -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, "")

18
litellm/rust_bridge/messages/detectors.py
View file

@ -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"})

96
tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py
View file

@ -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:

29
tests/test_litellm_rust/messages/test_dispatch.py
View file

@ -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

36
tests/unit/rust_bridge/host/routes/test_messages.py
View file

@ -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",
}

14
tests/unit/rust_bridge/test_detectors.py
View file

@ -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