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 a9ac07595f
commit ff2d7c985c
No known key found for this signature in database
18 changed files with 663 additions and 91 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_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<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,
@ -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<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()),
};
}

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

@ -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<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)] {
@ -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<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

@ -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<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);
@ -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<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| {
@ -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),
)

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

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

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

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

1
litellm/rust_bridge/catalog.py
View file

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

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

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

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

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

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

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

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

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

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

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

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

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