fix(rust-bridge): read secrets only through Python's secret manager

Rust routes on the Python bridge now resolve every secret through the
host module's read_secrets, which calls get_secret_str, so the proxy and
SDK have one resolver for a configured secret manager. A batch of names
is one host call.

The bypass that captured a Python manager's config and served reads from
a Rust backend is removed: ResolvedSecrets, the config snapshot and
callback, _SecretManagerRuntime.from_client, capture_secret_manager, and
the SecretManager/SecretManagerBinding settings snapshots.
_SecretManagerRuntime.from_config and the backend crates stay, for the
planned move of the Python managers onto native implementations.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-04 12:56:58 +09:00
parent cbe361d970
commit b74bd81eda
23 changed files with 153 additions and 1995 deletions

View file

@ -4331,7 +4331,6 @@ dependencies = [
name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"aws-sdk-secretsmanager",
"bytes",
"criterion",
"fancy-regex 0.19.2",
@ -4360,7 +4359,6 @@ dependencies = [
"litellm-router-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-secrets-aws",
"litellm-secrets-types",
"litellm-spend-clickhouse",
"litellm-storage-clickhouse",

View file

@ -66,7 +66,6 @@ tokio = { workspace = true, features = ["rt", "sync"] }
[dev-dependencies]
tracing.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-secrets-aws.workspace = true
serde.workspace = true
serde_with.workspace = true
criterion.workspace = true
@ -76,7 +75,6 @@ sqlx = { workspace = true, features = ["migrate"] }
tempfile.workspace = true
tokio-tungstenite.workspace = true
wiremock.workspace = true
aws-sdk-secretsmanager = "1.117.0"
[[bench]]
name = "serialization"

View file

@ -23,12 +23,6 @@ pub(crate) enum ProjectionDetail {
expected: &'static str,
actual: String,
},
#[error("{field}")]
InvalidField {
field: &'static str,
#[source]
source: serde_json::Error,
},
#[error("{group}.{name}: missing snapshot field")]
MissingSnapshotField {
group: &'static str,
@ -160,15 +154,6 @@ impl<'py> Field<'py> {
self.value.is(PyBool::new(self.value.py(), true))
}
pub(crate) fn schema_bool(&self) -> Result<bool, ProjectionError> {
if !self.value.is_instance_of::<PyBool>() {
return Err(ProjectionError::InternalSchemaFailure(
self.expected("a Boolean")?,
));
}
Ok(self.exact_true())
}
pub(crate) fn strict_string(&self) -> Result<String, ProjectionError> {
let Ok(value) = self.value.cast::<PyString>() else {
return Err(self.invalid("a string"));
@ -235,19 +220,6 @@ impl<'py> Field<'py> {
})
.collect()
}
pub(crate) fn optional_string_collection(
&self,
) -> Result<Option<Vec<String>>, ProjectionError> {
if self.value.is_none() {
return Ok(None);
}
self.string_collection().map(Some)
}
pub(crate) fn python_binding(&self) -> Option<Py<PyAny>> {
(!self.value.is_none()).then(|| self.value.clone().unbind())
}
}
#[cfg(test)]

View file

@ -375,7 +375,7 @@ mod tests {
Python::attach(|py| {
let raised = pyo3::exceptions::PyKeyError::new_err("vault down");
let identity = raised.value(py).clone().unbind();
let error = Error::from(crate::secrets::read_error(py, raised));
let error = Error::from(crate::secrets::external_error(py, raised));
let surfaced = route_failure(py, error, RequestMarker::Unmarked, |_| {
panic!("the raised exception must not be remapped")
})

View file

@ -13,10 +13,6 @@ pub(crate) enum PythonSettings {
UrlPolicy,
#[strum(serialize = "provider_defaults")]
ProviderDefaults,
#[strum(serialize = "secret_manager")]
SecretManager,
#[strum(serialize = "secret_manager_binding")]
SecretManagerBinding,
}
pub(crate) struct Snapshot<'py> {

View file

@ -1,508 +0,0 @@
use std::{future::Future, pin::Pin, sync::Arc};
use litellm_core_utils::settings::Lookup;
use litellm_host_python::{PythonContext, attach_blocking};
use litellm_secret_value::SecretValue;
use litellm_secrets::{
Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret,
};
use pyo3::{
exceptions::PyException,
prelude::*,
types::{PyDict, PyString},
};
use super::error::{external_error, read_error};
const HANDLER_MODULE: &str = "litellm.secret_managers.secret_manager_handler";
const DIAGNOSTICS_MODULE: &str = "litellm.rust_bridge.host.diagnostics";
/// A secret manager whose reads execute in Python: a custom manager, a legacy compatible
/// client, or a manually assigned SDK client.
pub(crate) struct PythonSecretManager {
client: Arc<PythonClient>,
context: PythonContext,
}
struct PythonClient {
client: Py<PyAny>,
system: Option<KeyManagementSystem>,
settings: Option<Py<PyAny>>,
}
impl PythonSecretManager {
pub(crate) fn new(
client: Py<PyAny>,
system: Option<KeyManagementSystem>,
settings: Option<Py<PyAny>>,
context: PythonContext,
) -> Self {
Self {
client: Arc::new(PythonClient {
client,
system,
settings,
}),
context,
}
}
}
impl PythonClient {
fn read(&self, py: Python<'_>, name: &str) -> PyResult<Option<String>> {
let client = self.client.bind(py);
let kwargs = PyDict::new(py);
kwargs.set_item("client", client)?;
kwargs.set_item("key_manager", self.system.map_or("local", python_name))?;
kwargs.set_item("secret_name", name)?;
match &self.settings {
Some(settings) => kwargs.set_item("key_management_settings", settings.bind(py))?,
None => kwargs.set_item("key_management_settings", py.None())?,
}
let result = py
.import(HANDLER_MODULE)?
.getattr("get_secret_from_manager")?
.call((), Some(&kwargs))?;
if result.is_instance_of::<PyString>() {
result.extract().map(Some)
} else {
Ok(None)
}
}
}
/// The `KeyManagementSystem` value as Python spells it.
fn python_name(system: KeyManagementSystem) -> &'static str {
system.into()
}
impl ExternalSecretManager for PythonSecretManager {
fn system(&self) -> KeyManagementSystem {
self.client.system.unwrap_or(KeyManagementSystem::Custom)
}
fn read_secret<'a>(
&'a self,
name: &'a str,
_settings: &'a KeyManagementSettings,
_environment: &'a (dyn Lookup + Send + Sync),
) -> Pin<Box<dyn Future<Output = Result<Option<Secret>, Error>> + Send + 'a>> {
let client = Arc::clone(&self.client);
let context = self.context.clone();
let name = name.to_owned();
Box::pin(async move {
match attach_blocking(context, move |py| match client.read(py, &name) {
Ok(value) => Ok(value.map(SecretValue::new).map(Secret::String)),
// `get_secret` answers a failed manager read from the process environment, but
// only for `Exception`: cancellation and other `BaseException`s propagate.
Err(error) if error.is_instance_of::<PyException>(py) => {
log_environment_fallback(py, &name, &error)
.map_err(|error| external_error(py, error))?;
Err(read_error(py, error))
}
Err(error) => Err(external_error(py, error)),
})
.await
{
Ok(result) => result,
Err(error) => Python::attach(|py| Err(external_error(py, error))),
}
})
}
}
fn log_environment_fallback(py: Python<'_>, name: &str, error: &PyErr) -> PyResult<()> {
py.import(DIAGNOSTICS_MODULE)?
.getattr("log_environment_fallback")?
.call1((name, error.value(py)))?;
Ok(())
}
#[cfg(test)]
#[allow(clippy::await_holding_lock)]
mod tests {
use std::sync::{Arc, Mutex, MutexGuard};
use litellm_secrets::{
FailurePolicy, KeyManagementSettings, KeyManagementSystem, OidcResolver, SecretManager,
SecretManagerState, SecretResolver,
};
use pyo3::{prelude::*, types::PyDict};
use rstest::rstest;
use litellm_host_python::PythonContext;
use super::{DIAGNOSTICS_MODULE, HANDLER_MODULE, PythonSecretManager, python_name};
use crate::secrets::python_error;
/// `sys.modules` is interpreter-global, so tests that install or rely on the handler module
/// cannot overlap with any other test on this list.
static HANDLER_LOCK: Mutex<()> = Mutex::new(());
fn handler_guard() -> MutexGuard<'static, ()> {
HANDLER_LOCK.lock().expect("handler lock poisoned")
}
/// A resolver over a Python manager whose reads raise `failure_type`, with the chained
/// exceptions Python attaches, and `fallback` as the process environment. The returned guard
/// keeps other module-mutating tests out for the lifetime of the returned resolver.
fn failing_resolver(
failure_type: &str,
fallback: Option<&'static str>,
) -> (SecretResolver, Py<PyDict>, MutexGuard<'static, ()>) {
let handler = handler_guard();
Python::initialize();
let (reader, locals) = Python::attach(|py| {
let locals = PyDict::new(py);
locals.set_item("failure_type", failure_type).unwrap();
py.run(
c"
import asyncio
failure = eval(failure_type)('secret manager failed')
cause = RuntimeError('original cause')
context = RuntimeError('original context')
failure.__cause__ = cause
failure.__context__ = context
class Manager:
def sync_read_secret(self, secret_name):
raise failure
manager = Manager()
import sys, types
for name in ('litellm', 'litellm.secret_managers'):
sys.modules.setdefault(name, types.ModuleType(name))
handler = sys.modules.setdefault('litellm.secret_managers.secret_manager_handler', types.ModuleType('litellm.secret_managers.secret_manager_handler'))
def get_secret_from_manager(**kwargs):
return kwargs['client'].sync_read_secret(kwargs['secret_name'])
handler.get_secret_from_manager = get_secret_from_manager
",
Some(&locals),
Some(&locals),
)
.unwrap();
let reader = PythonSecretManager::new(
locals.get_item("manager").unwrap().unwrap().unbind(),
None,
None,
PythonContext::capture(py).unwrap(),
);
(reader, locals.unbind())
});
let resolver = SecretResolver::new_python_compatible(
Arc::new(SecretManagerState::new(
SecretManager::External(Arc::new(reader)),
KeyManagementSettings::default(),
)),
Arc::new(move |_: &str| fallback.map(str::to_owned)),
OidcResolver::new(litellm_http::Client::for_test()),
)
.with_failure_policy(FailurePolicy::EnvironmentFallback);
(resolver, locals, handler)
}
#[rstest]
#[case::cancelled("asyncio.CancelledError", None)]
#[case::cancelled_with_fallback("asyncio.CancelledError", Some("environment-key"))]
#[case::keyboard_interrupt("KeyboardInterrupt", Some("environment-key"))]
#[tokio::test]
async fn base_exceptions_propagate_unchanged_even_with_environment_fallback(
#[case] failure_type: &str,
#[case] fallback: Option<&'static str>,
) {
let (resolver, locals, _handler) = failing_resolver(failure_type, fallback);
let error = resolver.get_secret("API_KEY", None).await.unwrap_err();
Python::attach(|py| {
let original = python_error(py, &error).unwrap();
let locals = locals.bind(py);
assert!(
original
.value(py)
.is(locals.get_item("failure").unwrap().unwrap())
);
for (attribute, name) in [("__cause__", "cause"), ("__context__", "context")] {
assert!(
original
.value(py)
.getattr(attribute)
.unwrap()
.is(locals.get_item(name).unwrap().unwrap())
);
}
assert!(original.traceback(py).is_some());
});
}
/// Installs a persistent `litellm._logging` stub whose `verbose_logger.error` records its
/// arguments, and returns those recorded for `name`.
fn logged_errors<'py>(py: Python<'py>, name: &str) -> Vec<Bound<'py, PyAny>> {
py.run(
c"
import sys, types
class Logger:
calls = []
def error(self, *args):
self.calls.append(args)
logging = types.ModuleType('litellm._logging')
logging.verbose_logger = Logger()
for module in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host'):
sys.modules.setdefault(module, types.ModuleType(module))
sys.modules.setdefault('litellm._logging', logging)
",
None,
None,
)
.unwrap();
pyo3::types::PyModule::from_code(
py,
&std::ffi::CString::new(include_str!(
"../../../../../litellm/rust_bridge/host/diagnostics.py"
))
.unwrap(),
c"diagnostics.py",
&std::ffi::CString::new(DIAGNOSTICS_MODULE).unwrap(),
)
.unwrap();
py.import("litellm._logging")
.unwrap()
.getattr("verbose_logger")
.unwrap()
.getattr("calls")
.unwrap()
.try_iter()
.unwrap()
.map(Result::unwrap)
.filter(|call| call.get_item(1).unwrap().extract::<String>().unwrap() == name)
.collect()
}
#[rstest]
#[case::value_error("ValueError", None, "FALLBACK_VALUE_ERROR")]
#[case::value_error_with_fallback(
"ValueError",
Some("environment-key"),
"FALLBACK_VALUE_ERROR_WITH_ENVIRONMENT"
)]
#[case::runtime_error_with_fallback(
"RuntimeError",
Some("environment-key"),
"FALLBACK_RUNTIME_ERROR_WITH_ENVIRONMENT"
)]
#[tokio::test]
async fn exceptions_are_logged_and_answered_from_the_environment(
#[case] failure_type: &str,
#[case] fallback: Option<&'static str>,
#[case] name: &str,
) {
let (resolver, _locals, _handler) = failing_resolver(failure_type, fallback);
Python::attach(|py| assert!(logged_errors(py, name).is_empty()));
let secret = resolver.get_secret(name, None).await.unwrap();
assert_eq!(
secret.map(|secret| match secret {
litellm_secrets::Secret::String(value) => value.expose().to_owned(),
other => panic!("unexpected secret {other:?}"),
}),
fallback.map(str::to_owned)
);
Python::attach(|py| {
let calls = logged_errors(py, name);
assert_eq!(calls.len(), 1);
assert!(
calls[0]
.get_item(3)
.unwrap()
.extract::<String>()
.unwrap()
.contains("sync_read_secret")
);
});
}
/// Installs a fake `get_secret_from_manager` that records its kwargs, runs `body`, and
/// removes the fake handler again; parent package stubs persist for concurrent tests.
/// Callers hold `handler_guard` before attaching so the GIL is never held while waiting on it.
fn with_fake_handler<'py>(py: Python<'py>, body: impl FnOnce(&Bound<'py, PyDict>)) {
let locals = PyDict::new(py);
py.run(
c"
import sys, types
previous_handler = sys.modules.get('litellm.secret_managers.secret_manager_handler')
calls = []
def get_secret_from_manager(**kwargs):
calls.append(kwargs)
return 'handled-' + kwargs['secret_name']
handler = types.ModuleType('litellm.secret_managers.secret_manager_handler')
handler.get_secret_from_manager = get_secret_from_manager
for name in ('litellm', 'litellm.secret_managers'):
sys.modules.setdefault(name, types.ModuleType(name))
sys.modules['litellm.secret_managers.secret_manager_handler'] = handler
",
Some(&locals),
Some(&locals),
)
.unwrap();
body(&locals);
py.run(
c"
if previous_handler is None:
sys.modules.pop('litellm.secret_managers.secret_manager_handler', None)
else:
sys.modules['litellm.secret_managers.secret_manager_handler'] = previous_handler
",
Some(&locals),
Some(&locals),
)
.unwrap();
}
#[rstest]
#[case("None")]
#[case("True")]
#[case("123")]
#[case("{'key': 'value'}")]
fn nonstring_results_are_absent_without_a_read_failure(#[case] expression: &str) {
let _handler = handler_guard();
Python::initialize();
Python::attach(|py| {
with_fake_handler(py, |locals| {
locals.set_item("expression", expression).unwrap();
py.run(
c"handler.get_secret_from_manager = lambda **kwargs: eval(expression)",
Some(locals),
Some(locals),
)
.unwrap();
let reader = PythonSecretManager::new(
py.None(),
None,
None,
PythonContext::capture(py).unwrap(),
);
assert_eq!(reader.client.read(py, "KEY").unwrap(), None);
});
});
}
#[rstest]
#[case::google_kms(KeyManagementSystem::GoogleKms)]
#[case::azure_key_vault(KeyManagementSystem::AzureKeyVault)]
#[case::aws_secret_manager(KeyManagementSystem::AwsSecretManager)]
#[case::google_secret_manager(KeyManagementSystem::GoogleSecretManager)]
#[case::hashicorp_vault(KeyManagementSystem::HashicorpVault)]
#[case::cyberark(KeyManagementSystem::Cyberark)]
#[case::local(KeyManagementSystem::Local)]
#[case::aws_kms(KeyManagementSystem::AwsKms)]
#[case::custom(KeyManagementSystem::Custom)]
fn python_names_round_trip_through_serde(#[case] system: KeyManagementSystem) {
assert_eq!(
serde_json::to_value(system).unwrap(),
serde_json::Value::String(python_name(system).to_owned())
);
}
#[test]
fn configured_systems_dispatch_through_the_python_handler_with_the_original_settings() {
let _handler = handler_guard();
Python::initialize();
Python::attach(|py| {
with_fake_handler(py, |locals| {
let client = py.eval(c"object()", None, None).unwrap();
let settings = py.eval(c"object()", None, None).unwrap();
let reader = PythonSecretManager::new(
client.clone().unbind(),
Some(KeyManagementSystem::AzureKeyVault),
Some(settings.clone().unbind()),
PythonContext::capture(py).unwrap(),
);
assert_eq!(
reader.client.read(py, "API_KEY").unwrap().as_deref(),
Some("handled-API_KEY")
);
assert!(py.import(HANDLER_MODULE).is_ok());
let calls = locals.get_item("calls").unwrap().unwrap();
let call = calls.get_item(0).unwrap().cast_into::<PyDict>().unwrap();
assert!(call.get_item("client").unwrap().unwrap().is(&client));
assert!(
call.get_item("key_management_settings")
.unwrap()
.unwrap()
.is(&settings)
);
assert_eq!(
call.get_item("key_manager")
.unwrap()
.unwrap()
.extract::<String>()
.unwrap(),
"azure_key_vault"
);
assert_eq!(
call.get_item("secret_name")
.unwrap()
.unwrap()
.extract::<String>()
.unwrap(),
"API_KEY"
);
});
});
}
#[rstest]
#[case::manually_assigned(None, "local")]
#[case::custom(Some(KeyManagementSystem::Custom), "custom")]
fn direct_readers_dispatch_through_the_python_handler_like_get_secret(
#[case] system: Option<KeyManagementSystem>,
#[case] key_manager: &str,
) {
let _handler = handler_guard();
Python::initialize();
Python::attach(|py| {
with_fake_handler(py, |locals| {
py.run(
c"
class Manager:
def __init__(self):
self.names = []
def sync_read_secret(self, secret_name, optional_params=None, timeout=None):
self.names.append(secret_name)
return 'direct-' + secret_name
manager = Manager()
",
Some(locals),
Some(locals),
)
.unwrap();
let manager = locals.get_item("manager").unwrap().unwrap();
let reader = PythonSecretManager::new(
manager.clone().unbind(),
system,
None,
PythonContext::capture(py).unwrap(),
);
assert_eq!(
reader.client.read(py, "API_KEY").unwrap().as_deref(),
Some("handled-API_KEY")
);
assert_eq!(
manager
.getattr("names")
.unwrap()
.extract::<Vec<String>>()
.unwrap(),
Vec::<String>::new()
);
let calls = locals.get_item("calls").unwrap().unwrap();
let call = calls.get_item(0).unwrap().cast_into::<PyDict>().unwrap();
assert!(call.get_item("client").unwrap().unwrap().is(&manager));
assert_eq!(
call.get_item("key_manager")
.unwrap()
.unwrap()
.extract::<String>()
.unwrap(),
key_manager
);
});
});
}
}

View file

@ -1,385 +0,0 @@
use std::sync::Arc;
use litellm_secret_value::SecretValue;
use litellm_secrets::{SecretManager, SecretManagerState};
use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem};
use pyo3::prelude::*;
use serde_json::Value;
use litellm_host_python::PythonContext;
use super::callback::PythonSecretManager;
use crate::{
coercion::{Field, FieldSpec, ProjectionError},
python_settings::{PythonSettings, Snapshot},
};
const SYSTEM: FieldSpec<Option<KeyManagementSystem>> =
FieldSpec::new("system", parse_optional_system);
const ACCESS_MODE: FieldSpec<AccessMode> = FieldSpec::new("access_mode", parse_access_mode);
const HOSTED_KEYS: FieldSpec<Option<Vec<String>>> =
FieldSpec::new("hosted_keys", |field| field.optional_string_collection());
const STORE_VIRTUAL_KEYS: FieldSpec<bool> =
FieldSpec::new("store_virtual_keys", |field| field.truthy());
const PREFIX_FOR_STORED_VIRTUAL_KEYS: FieldSpec<String> =
FieldSpec::new("prefix_for_stored_virtual_keys", |field| {
field.strict_string()
});
const PRIMARY_SECRET_NAME: FieldSpec<Option<String>> =
FieldSpec::new("primary_secret_name", |field| field.falsy_optional_string());
const KMS_KEY_ID: FieldSpec<Option<String>> =
FieldSpec::new("kms_key_id", |field| field.falsy_optional_string());
const CUSTOM_SECRET_MANAGER: FieldSpec<Option<String>> =
FieldSpec::new("custom_secret_manager", |field| {
field.falsy_optional_string()
});
const AWS_REGION_NAME: FieldSpec<Option<String>> =
FieldSpec::new("aws_region_name", |field| field.falsy_optional_string());
const AWS_ROLE_NAME: FieldSpec<Option<String>> =
FieldSpec::new("aws_role_name", |field| field.falsy_optional_string());
const AWS_SESSION_NAME: FieldSpec<Option<String>> =
FieldSpec::new("aws_session_name", |field| field.falsy_optional_string());
const AWS_EXTERNAL_ID: FieldSpec<Option<String>> =
FieldSpec::new("aws_external_id", |field| field.falsy_optional_string());
const AWS_PROFILE_NAME: FieldSpec<Option<String>> =
FieldSpec::new("aws_profile_name", |field| field.falsy_optional_string());
const AWS_WEB_IDENTITY_TOKEN: FieldSpec<Option<String>> =
FieldSpec::new("aws_web_identity_token", |field| {
field.falsy_optional_string()
});
const AWS_STS_ENDPOINT: FieldSpec<Option<String>> =
FieldSpec::new("aws_sts_endpoint", |field| field.falsy_optional_string());
const REPLICA_REGIONS: FieldSpec<Option<Vec<String>>> =
FieldSpec::new("replica_regions", |field| {
field.optional_string_collection()
});
const CLIENT: FieldSpec<Option<Py<PyAny>>> =
FieldSpec::new("client", |field| Ok(field.python_binding()));
const SETTINGS_OBJECT: FieldSpec<Option<Py<PyAny>>> =
FieldSpec::new("settings_object", |field| Ok(field.python_binding()));
/// `litellm.secret_manager_client` as the bridge classifies it.
pub(crate) enum SecretManagerClient {
/// `None`: reads come from the process environment.
Local,
/// A custom manager, legacy compatible client, or manually assigned SDK client that keeps
/// executing in Python.
PythonCallback(Py<PyAny>),
Native(Box<SecretManager>),
}
impl std::fmt::Debug for SecretManagerClient {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(match self {
Self::Local => "Local",
Self::Native(_) => "Native",
Self::PythonCallback(_) => "PythonCallback",
})
}
}
/// One operation-local capture of the secret manager globals, taken while attached to Python.
#[derive(Debug)]
pub(crate) struct SecretManagerSnapshot {
pub(crate) client: SecretManagerClient,
pub(crate) system: Option<KeyManagementSystem>,
/// Typed settings that drive native routing: access mode and hosted keys.
pub(crate) settings: KeyManagementSettings,
/// The original `KeyManagementSettings` object, handed back to Python callbacks unchanged.
pub(crate) settings_object: Option<Py<PyAny>>,
}
impl SecretManagerSnapshot {
pub(crate) fn into_state(self, context: PythonContext) -> Arc<SecretManagerState> {
match self.client {
SecretManagerClient::Native(backend) => {
Arc::new(SecretManagerState::new(*backend, self.settings))
}
SecretManagerClient::Local => Arc::new(SecretManagerState::default()),
SecretManagerClient::PythonCallback(client) => Arc::new(SecretManagerState::new(
SecretManager::External(Arc::new(PythonSecretManager::new(
client,
self.system,
self.settings_object,
context,
))),
self.settings,
)),
}
}
}
/// Reads and projects the secret manager settings group in one attached operation.
pub(crate) fn read(py: Python<'_>) -> PyResult<SecretManagerSnapshot> {
let snapshot = project(&PythonSettings::SecretManagerBinding.read(py)?)?;
let SecretManagerClient::PythonCallback(client) = &snapshot.client else {
return Ok(snapshot);
};
if matches!(
snapshot.system,
Some(KeyManagementSystem::Custom | KeyManagementSystem::Local)
) {
return Ok(snapshot);
}
let Some(native) = super::runtime::NativeSecretManager::from_client(client.bind(py))? else {
return Ok(snapshot);
};
let backend = native.borrow(py).backend()?;
if snapshot
.system
.is_some_and(|system| system != backend.system())
{
return Err(super::ManagerError::SystemMismatch.into());
}
Ok(SecretManagerSnapshot {
client: SecretManagerClient::Native(Box::new(backend)),
..snapshot
})
}
pub(crate) fn project(snapshot: &Snapshot<'_>) -> Result<SecretManagerSnapshot, ProjectionError> {
let system = snapshot.read(&SYSTEM)?;
let access_mode = snapshot.read(&ACCESS_MODE)?;
let settings = KeyManagementSettings {
hosted_keys: snapshot.read(&HOSTED_KEYS)?,
store_virtual_keys: Some(snapshot.read(&STORE_VIRTUAL_KEYS)?),
prefix_for_stored_virtual_keys: snapshot.read(&PREFIX_FOR_STORED_VIRTUAL_KEYS)?,
access_mode,
primary_secret_name: snapshot.read(&PRIMARY_SECRET_NAME)?,
kms_key_id: snapshot.read(&KMS_KEY_ID)?,
custom_secret_manager: snapshot.read(&CUSTOM_SECRET_MANAGER)?,
aws_region_name: snapshot.read(&AWS_REGION_NAME)?,
aws_role_name: snapshot.read(&AWS_ROLE_NAME)?,
aws_session_name: snapshot.read(&AWS_SESSION_NAME)?,
aws_external_id: snapshot.read(&AWS_EXTERNAL_ID)?.map(SecretValue::new),
aws_profile_name: snapshot.read(&AWS_PROFILE_NAME)?,
aws_web_identity_token: snapshot
.read(&AWS_WEB_IDENTITY_TOKEN)?
.map(SecretValue::new),
aws_sts_endpoint: snapshot.read(&AWS_STS_ENDPOINT)?,
replica_regions: snapshot.read(&REPLICA_REGIONS)?,
..KeyManagementSettings::default()
};
let client = match snapshot.read(&CLIENT)? {
None => SecretManagerClient::Local,
Some(client) => SecretManagerClient::PythonCallback(client),
};
Ok(SecretManagerSnapshot {
client,
system,
settings,
settings_object: snapshot.read(&SETTINGS_OBJECT)?,
})
}
fn parse_optional_system(
field: &Field<'_>,
) -> Result<Option<KeyManagementSystem>, ProjectionError> {
let Some(value) = field.falsy_optional_string()? else {
return Ok(None);
};
serde_json::from_value(Value::String(value))
.map(Some)
.map_err(|error| {
ProjectionError::InvalidConfiguration(crate::coercion::ProjectionDetail::InvalidField {
field: "secret manager system",
source: error,
})
})
}
fn parse_access_mode(field: &Field<'_>) -> Result<AccessMode, ProjectionError> {
let value = field.strict_string()?;
serde_json::from_value(Value::String(value)).map_err(|error| {
ProjectionError::InvalidConfiguration(crate::coercion::ProjectionDetail::InvalidField {
field: "secret manager access mode",
source: error,
})
})
}
#[cfg(test)]
mod tests {
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use super::{SecretManagerClient, project};
use crate::python_settings::PythonSettings;
fn snapshot<'py>(
py: Python<'py>,
system: &str,
access_mode: &str,
store_virtual_keys: Bound<'py, PyAny>,
hosted_keys: Bound<'py, PyAny>,
) -> crate::python_settings::Snapshot<'py> {
snapshot_with_client(
py,
system,
access_mode,
store_virtual_keys,
hosted_keys,
py.None().into_bound(py),
)
}
fn snapshot_with_client<'py>(
py: Python<'py>,
system: &str,
access_mode: &str,
store_virtual_keys: Bound<'py, PyAny>,
hosted_keys: Bound<'py, PyAny>,
client: Bound<'py, PyAny>,
) -> crate::python_settings::Snapshot<'py> {
let locals = PyDict::new(py);
locals.set_item("client", client).unwrap();
locals.set_item("system", system).unwrap();
locals.set_item("access_mode", access_mode).unwrap();
locals
.set_item("store_virtual_keys", store_virtual_keys)
.unwrap();
locals.set_item("hosted_keys", hosted_keys).unwrap();
py.run(
cr#"
from dataclasses import dataclass
from types import SimpleNamespace
@dataclass(frozen=True, slots=True)
class SecretManager:
system: object
access_mode: object
hosted_keys: object
primary_secret_name: object
store_virtual_keys: object
prefix_for_stored_virtual_keys: object
kms_key_id: object
custom_secret_manager: object
aws_region_name: object
aws_role_name: object
aws_session_name: object
aws_external_id: object
aws_profile_name: object
aws_web_identity_token: object
aws_sts_endpoint: object
replica_regions: object
client: object
settings_object: object
root = SimpleNamespace(secret_manager=SecretManager(
system=system,
access_mode=access_mode,
hosted_keys=hosted_keys,
primary_secret_name=None,
store_virtual_keys=store_virtual_keys,
prefix_for_stored_virtual_keys="litellm/",
kms_key_id=None,
custom_secret_manager=None,
aws_region_name=None,
aws_role_name=None,
aws_session_name=None,
aws_external_id=None,
aws_profile_name=None,
aws_web_identity_token=None,
aws_sts_endpoint=None,
replica_regions=None,
client=client,
settings_object=None,
))
"#,
Some(&locals),
Some(&locals),
)
.unwrap();
PythonSettings::SecretManagerBinding.snapshot(
locals
.get_item("root")
.unwrap()
.unwrap()
.getattr("secret_manager")
.unwrap(),
)
}
#[rstest::rstest]
#[case::string_true(Some("true"), false, true)]
#[case::string_one(Some("1"), false, true)]
#[case::true_value(None, true, true)]
#[case::false_value(None, false, false)]
#[case::string_false(Some("false"), false, true)]
fn python_compatible_boolean_coercion(
#[case] string_value: Option<&str>,
#[case] bool_value: bool,
#[case] expected: bool,
) {
Python::initialize();
Python::attach(|py| {
let store_virtual_keys = match string_value {
Some(value) => value.into_pyobject(py).unwrap().into_any(),
None => bool_value.into_pyobject(py).unwrap().to_owned().into_any(),
};
let hosted_keys = PyTuple::new(py, ["ONE"]).unwrap().into_any();
let projected = project(&snapshot(
py,
"local",
"read_only",
store_virtual_keys,
hosted_keys,
))
.unwrap();
assert_eq!(projected.settings.store_virtual_keys, Some(expected));
});
}
#[test]
fn unknown_system_is_rejected() {
Python::initialize();
Python::attach(|py| {
let error = project(&snapshot(
py,
"unknown",
"read_only",
false.into_pyobject(py).unwrap().to_owned().into_any(),
PyTuple::empty(py).into_any(),
))
.unwrap_err();
let error: PyErr = error.into();
assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
});
}
#[test]
fn client_identity_selects_local_or_python_callback() {
Python::initialize();
Python::attach(|py| {
let falsy = false.into_pyobject(py).unwrap().to_owned().into_any();
let local = project(&snapshot(
py,
"local",
"read_only",
falsy.clone(),
PyTuple::empty(py).into_any(),
))
.unwrap();
assert!(matches!(local.client, SecretManagerClient::Local));
assert!(local.settings_object.is_none());
let manager = py.eval(c"object()", None, None).unwrap();
let custom = project(&snapshot_with_client(
py,
"custom",
"read_only",
falsy,
PyTuple::empty(py).into_any(),
manager.clone(),
))
.unwrap();
let SecretManagerClient::PythonCallback(client) = custom.client else {
panic!("a live client must stay a Python callback");
};
assert!(client.bind(py).is(&manager));
});
}
}

View file

@ -16,8 +16,6 @@ pub(crate) enum ManagerError {
UnknownSystem(#[source] serde_json::Error),
#[error("invalid secret manager settings")]
Settings(#[source] litellm_error::Source),
#[error("native secret manager system does not match configuration")]
SystemMismatch,
#[error("secret manager does not support provider reads")]
ProviderReads,
#[error("invalid timeout")]
@ -49,7 +47,6 @@ impl From<ManagerError> for PyErr {
ManagerError::Forked => PyRuntimeError::new_err(message),
ManagerError::UnknownSystem(_)
| ManagerError::Settings(_)
| ManagerError::SystemMismatch
| ManagerError::ProviderReads
| ManagerError::Timeout(_)
| ManagerError::CyberarkSecretMissing(_)
@ -72,18 +69,12 @@ impl fmt::Debug for PythonSecretError {
}
}
pub(super) fn external_error(py: Python<'_>, error: PyErr) -> Error {
pub(crate) fn external_error(py: Python<'_>, error: PyErr) -> Error {
Error::ExternalManager(litellm_error::Source::new(PythonSecretError(
error.into_value(py),
)))
}
pub(crate) fn read_error(py: Python<'_>, error: PyErr) -> Error {
Error::ExternalRead(litellm_error::Source::new(PythonSecretError(
error.into_value(py),
)))
}
/// The exception a Python secret manager raised anywhere in the chain, so a route surfaces it
/// unchanged instead of wrapping it.
pub(crate) fn python_error(

View file

@ -1,46 +1,34 @@
pub(crate) mod callback;
pub(crate) mod config;
mod error;
mod mutation;
mod operations;
mod provider;
mod python;
pub(crate) mod resolved;
pub(crate) mod runtime;
mod vault;
use std::sync::Arc;
#[cfg(test)]
pub(crate) use error::read_error;
pub(crate) use error::external_error;
pub(crate) use error::{ManagerError, python_error};
use litellm_secrets::source::{EnvironmentSecrets, SecretSource};
use pyo3::prelude::*;
use python::PythonSecrets;
use resolved::ResolvedSecrets;
use crate::{coercion::FieldSpec, python_settings::PythonSettings};
use crate::python_settings::missing_module;
const NATIVE: FieldSpec<bool> = FieldSpec::new("native", |field| field.schema_bool());
/// Where a Rust route reads provider secrets from. Python's `get_secret_str` when the settings
/// projection reports no configured secret manager, otherwise the resolved manager: the native
/// backend when the configured client captures one, the Python callback when it does not. A bare
/// extension module without the litellm package reads the process environment.
/// Where a Rust route reads provider secrets from: Python's `get_secret_str`, the one resolver for
/// the configured secret manager. A bare extension module without the litellm package reads the
/// process environment.
pub(crate) fn source(py: Python<'_>) -> PyResult<Arc<dyn SecretSource>> {
let Some(snapshot) = PythonSettings::SecretManager.read_or_unset(py)? else {
return Ok(Arc::new(EnvironmentSecrets::python_compatible(
crate::http::host_client(py, litellm_http::ClientVariant::Provider)?,
)));
};
if snapshot.read(&NATIVE)? {
let context = litellm_host_python::PythonContext::capture(py)?;
let client = crate::http::host_client(py, litellm_http::ClientVariant::Provider)?;
return Ok(Arc::new(ResolvedSecrets::new(
config::read(py)?,
context,
client,
)));
match PythonSecrets::new(py) {
Ok(secrets) => Ok(Arc::new(secrets)),
Err(error) if missing_module(py, &error, "litellm")? => Ok(Arc::new(
EnvironmentSecrets::python_compatible(crate::http::host_client(
py,
litellm_http::ClientVariant::Provider,
)?),
)),
Err(error) => Err(error),
}
Ok(Arc::new(PythonSecrets::new(py)?))
}

View file

@ -1,20 +1,25 @@
use std::sync::Arc;
use std::{collections::HashMap, sync::Arc};
use futures_util::future::BoxFuture;
use litellm_core_utils::settings::Lookup;
use litellm_host_python::{PythonContext, attach_blocking};
use litellm_secret_value::SecretValue;
use litellm_secrets::{Error, source::SecretSource};
use litellm_secrets::{
Error,
source::{SecretSource, Secrets},
};
use pyo3::prelude::*;
use pyo3::{prelude::*, types::PyTuple};
use super::error::external_error;
const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager";
/// Reads each secret through the host module's `read_secret_str`, so the configured manager, the key
/// management settings and the environment fallback behave exactly as they do in Python.
/// Reads secrets through the host module's `read_secrets`, so the configured manager, the key
/// management settings and the environment fallback behave exactly as they do in Python. A batch of
/// names is one host call.
pub(super) struct PythonSecrets {
get_secret_str: Arc<Py<PyAny>>,
read_secrets: Arc<Py<PyAny>>,
context: PythonContext,
}
@ -22,18 +27,42 @@ impl PythonSecrets {
pub(super) fn new(py: Python<'_>) -> PyResult<Self> {
Ok(Self::reading_with(
py.import(SECRET_MANAGER_MODULE)?
.getattr("read_secret_str")?
.getattr("read_secrets")?
.unbind(),
PythonContext::capture(py)?,
))
}
fn reading_with(get_secret_str: Py<PyAny>, context: PythonContext) -> Self {
fn reading_with(read_secrets: Py<PyAny>, context: PythonContext) -> Self {
Self {
get_secret_str: Arc::new(get_secret_str),
read_secrets: Arc::new(read_secrets),
context,
}
}
async fn read(&self, names: Vec<String>) -> Result<Vec<Option<SecretValue>>, Error> {
let read_secrets = Arc::clone(&self.read_secrets);
match attach_blocking(self.context.clone(), move |py| {
PyTuple::new(py, names)
.and_then(|names| read_secrets.bind(py).call1((names,)))
.and_then(|values| values.extract::<Vec<Option<String>>>())
.map(|values| values.into_iter().map(|value| value.map(SecretValue::new)).collect())
.map_err(|error| external_error(py, error))
})
.await
{
Ok(result) => result,
Err(error) => Python::attach(|py| Err(external_error(py, error))),
}
}
}
struct ReadSecrets(HashMap<String, SecretValue>);
impl Lookup for ReadSecrets {
fn get(&self, name: &str) -> Option<String> {
self.0.get(name).map(|value| value.expose().to_owned())
}
}
impl SecretSource for PythonSecrets {
@ -41,23 +70,27 @@ impl SecretSource for PythonSecrets {
&'a self,
name: &'a str,
) -> BoxFuture<'a, Result<Option<SecretValue>, Error>> {
let get_secret_str = Arc::clone(&self.get_secret_str);
let context = self.context.clone();
let name = name.to_owned();
Box::pin(async move {
match attach_blocking(context, move |py| {
get_secret_str
.bind(py)
.call1((name,))
.and_then(|value| value.extract::<Option<String>>())
.map(|value| value.map(SecretValue::new))
.map_err(|error| external_error(py, error))
})
.await
{
Ok(result) => result,
Err(error) => Python::attach(|py| Err(external_error(py, error))),
}
Ok(self
.read(vec![name.to_owned()])
.await?
.into_iter()
.next()
.flatten())
})
}
fn resolve<'a>(&'a self, names: &'a [&str]) -> BoxFuture<'a, Result<Secrets, Error>> {
Box::pin(async move {
let values = self
.read(names.iter().map(|name| (*name).to_owned()).collect())
.await?;
let found = names
.iter()
.zip(values)
.filter_map(|(name, value)| Some(((*name).to_owned(), value?)))
.collect();
Ok(Arc::new(ReadSecrets(found)) as Secrets)
})
}
}
@ -84,14 +117,17 @@ import threading
read_on = None
request_var = contextvars.ContextVar('request_var', default=None)
seen_context_values = []
calls = []
raised = KeyboardInterrupt('secret manager stopped')
def get_secret_str(name):
values = {'MISTRAL_API_KEY': 'vault-key', 'MISTRAL_API_BASE': 'https://mistral.internal'}
def read_secrets(names):
global read_on
read_on = threading.get_ident()
seen_context_values.append(request_var.get())
if name == 'RAISING':
calls.append(names)
if 'RAISING' in names:
raise raised
return {'MISTRAL_API_KEY': 'vault-key'}.get(name)
return tuple(values.get(name) for name in names)
",
Some(&namespace),
None,
@ -113,7 +149,7 @@ def get_secret_str(name):
.unwrap();
(
namespace
.get_item("get_secret_str")
.get_item("read_secrets")
.unwrap()
.unwrap()
.unbind(),
@ -141,7 +177,7 @@ def get_secret_str(name):
#[case::found("MISTRAL_API_KEY", Some("vault-key"))]
#[case::missing("OTHER", None)]
#[tokio::test]
async fn returns_what_get_secret_str_returns(
async fn returns_what_read_secrets_returns(
secrets: (PythonSecrets, Py<PyDict>),
#[case] name: &str,
#[case] expected: Option<&str>,
@ -151,12 +187,43 @@ def get_secret_str(name):
assert_eq!(value.as_ref().map(|value| value.expose()), expected);
}
#[rstest]
#[tokio::test]
async fn resolve_reads_every_name_in_one_python_call(secrets: (PythonSecrets, Py<PyDict>)) {
let resolved = secrets
.0
.resolve(&["MISTRAL_API_KEY", "OTHER", "MISTRAL_API_BASE"])
.await
.unwrap();
assert_eq!(resolved.get("MISTRAL_API_KEY").as_deref(), Some("vault-key"));
assert_eq!(resolved.get("OTHER"), None);
assert_eq!(
resolved.get("MISTRAL_API_BASE").as_deref(),
Some("https://mistral.internal")
);
let calls: Vec<Vec<String>> = Python::attach(|py| global(&secrets.1, py, "calls"));
assert_eq!(
calls,
vec![vec![
"MISTRAL_API_KEY".to_owned(),
"OTHER".to_owned(),
"MISTRAL_API_BASE".to_owned()
]]
);
}
#[rstest]
#[tokio::test]
async fn exceptions_surface_as_the_original_python_object(
secrets: (PythonSecrets, Py<PyDict>),
) {
let error = secrets.0.get_secret_str("RAISING").await.unwrap_err();
let error = secrets
.0
.resolve(&["MISTRAL_API_KEY", "RAISING"])
.await
.err()
.unwrap();
Python::attach(|py| {
let surfaced = python_error(py, &error).expect("the Python exception is preserved");
@ -186,7 +253,7 @@ def get_secret_str(name):
#[rstest]
#[tokio::test]
async fn reads_see_the_callers_contextvars(secrets: (PythonSecrets, Py<PyDict>)) {
secrets.0.get_secret_str("MISTRAL_API_KEY").await.unwrap();
secrets.0.resolve(&["MISTRAL_API_KEY"]).await.unwrap();
let seen: Vec<String> = Python::attach(|py| global(&secrets.1, py, "seen_context_values"));
assert_eq!(seen, vec!["request-value".to_owned()]);

View file

@ -1,277 +0,0 @@
use std::sync::Arc;
use futures_util::future::BoxFuture;
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_host_python::PythonContext;
use litellm_http::Client;
use litellm_secret_value::SecretValue;
use litellm_secrets::source::SecretSource;
use litellm_secrets::{Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver};
use super::config::SecretManagerSnapshot;
pub(crate) struct ResolvedSecrets {
resolver: SecretResolver,
}
impl ResolvedSecrets {
pub(crate) fn new(
snapshot: SecretManagerSnapshot,
context: PythonContext,
client: Client,
) -> Self {
Self::from_state(snapshot.into_state(context), client)
}
fn from_state(state: Arc<SecretManagerState>, client: Client) -> Self {
Self {
resolver: SecretResolver::new_python_compatible(
state,
Arc::new(ProcessEnvironment),
OidcResolver::new(client),
)
.with_failure_policy(FailurePolicy::EnvironmentFallback),
}
}
}
impl SecretSource for ResolvedSecrets {
fn get_secret_str<'a>(
&'a self,
name: &'a str,
) -> BoxFuture<'a, Result<Option<SecretValue>, Error>> {
Box::pin(self.resolver.get_secret_str(name, None))
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aws_sdk_secretsmanager::Client;
use aws_sdk_secretsmanager::config::{
BehaviorVersion, Credentials, Region, retry::RetryConfig,
};
use litellm_secrets::{AccessMode, KeyManagementSettings, SecretManager, SecretManagerState};
use litellm_secrets_aws::AwsSecretsManagerV2;
use serde_json::json;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_partial_json, header},
};
use super::ResolvedSecrets;
use litellm_secrets::source::SecretSource;
fn state(server: &MockServer, settings: KeyManagementSettings) -> Arc<SecretManagerState> {
let client = Client::from_conf(
aws_sdk_secretsmanager::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new("us-east-1"))
.credentials_provider(Credentials::new("test", "test", None, None, "test"))
.endpoint_url(server.uri())
.retry_config(RetryConfig::disabled())
.build(),
);
Arc::new(SecretManagerState::new(
SecretManager::AwsSecretsManagerV2(AwsSecretsManagerV2::new(
client,
(&settings).into(),
)),
settings,
))
}
async fn resolve(state: Arc<SecretManagerState>, name: &'static str) -> Option<String> {
ResolvedSecrets::from_state(state, litellm_http::Client::for_test())
.resolve(&[name])
.await
.unwrap()
.get(name)
}
#[tokio::test]
async fn hosted_key_miss_falls_back_to_environment() {
let name = "LITELLM_RUST_BRIDGE_HOSTED_KEY_MISS";
unsafe { std::env::set_var(name, "env-key") };
let server = MockServer::start().await;
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
.and(body_partial_json(json!({"SecretId": name})))
.respond_with(
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
)
.expect(0)
.mount(&server)
.await;
let result = resolve(
state(
&server,
KeyManagementSettings {
hosted_keys: Some(vec!["OTHER".into()]),
..Default::default()
},
),
name,
)
.await;
unsafe { std::env::remove_var(name) };
assert_eq!(result.as_deref(), Some("env-key"));
assert_eq!(server.received_requests().await.unwrap().len(), 0);
}
#[tokio::test]
async fn aws_read_failure_preserves_absence_without_environment_fallback() {
let name = "LITELLM_RUST_BRIDGE_MANAGER_FAILURE";
unsafe { std::env::set_var(name, "env-key") };
let server = MockServer::start().await;
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
.respond_with(ResponseTemplate::new(500))
.expect(2)
.mount(&server)
.await;
let result = resolve(state(&server, KeyManagementSettings::default()), name).await;
let missing = resolve(
state(&server, KeyManagementSettings::default()),
"LITELLM_RUST_BRIDGE_MANAGER_FAILURE_MISSING",
)
.await;
unsafe { std::env::remove_var(name) };
assert_eq!(result, None);
assert_eq!(missing, None);
}
#[rstest::rstest]
#[case::capitalized_true("True")]
#[case::parenthesized_false("(False)")]
#[tokio::test]
async fn boolean_manager_values_are_absent_like_get_secret_str(#[case] value: &str) {
let name = "LITELLM_RUST_BRIDGE_BOOLEAN_VALUE";
let server = MockServer::start().await;
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"SecretString": value})))
.expect(1)
.mount(&server)
.await;
assert_eq!(
resolve(state(&server, KeyManagementSettings::default()), name).await,
None
);
}
#[tokio::test]
async fn undeclared_names_are_read_from_the_manager() {
let declared = "LITELLM_RUST_BRIDGE_DECLARED";
let undeclared = "LITELLM_RUST_BRIDGE_UNDECLARED_MANAGED";
unsafe { std::env::set_var(undeclared, "env-key") };
let server = MockServer::start().await;
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
.and(body_partial_json(json!({"SecretId": declared})))
.respond_with(
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "declared-key"})),
)
.mount(&server)
.await;
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
.and(body_partial_json(json!({"SecretId": undeclared})))
.respond_with(
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
)
.expect(1)
.mount(&server)
.await;
let source = ResolvedSecrets::from_state(
state(&server, KeyManagementSettings::default()),
litellm_http::Client::for_test(),
);
let snapshot = source.resolve(&[declared]).await.unwrap();
assert_eq!(snapshot.get(undeclared), None);
let result = source
.get_secret_str(undeclared)
.await
.unwrap()
.map(|value| value.expose().to_owned());
unsafe { std::env::remove_var(undeclared) };
assert_eq!(result.as_deref(), Some("manager-key"));
}
#[tokio::test]
async fn write_only_mode_never_consults_the_manager() {
let name = "LITELLM_RUST_BRIDGE_WRITE_ONLY";
unsafe { std::env::set_var(name, "env-key") };
let server = MockServer::start().await;
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
.respond_with(
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
)
.expect(0)
.mount(&server)
.await;
let result = resolve(
state(
&server,
KeyManagementSettings {
access_mode: AccessMode::WriteOnly,
..Default::default()
},
),
name,
)
.await;
unsafe { std::env::remove_var(name) };
assert_eq!(result.as_deref(), Some("env-key"));
assert_eq!(server.received_requests().await.unwrap().len(), 0);
}
#[tokio::test]
async fn read_only_mode_resolves_from_the_manager() {
let name = "LITELLM_RUST_BRIDGE_READ_ONLY";
let server = MockServer::start().await;
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
.and(body_partial_json(json!({"SecretId": name})))
.respond_with(
ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})),
)
.expect(1)
.mount(&server)
.await;
assert_eq!(
resolve(state(&server, KeyManagementSettings::default()), name)
.await
.as_deref(),
Some("manager-key")
);
assert_eq!(server.received_requests().await.unwrap().len(), 1);
}
#[tokio::test]
async fn oidc_failures_are_not_converted_to_missing_secrets() {
let result = ResolvedSecrets::from_state(
Arc::new(SecretManagerState::default()),
litellm_http::Client::for_test(),
)
.resolve(&["oidc/"])
.await;
assert!(matches!(result, Err(litellm_secrets::Error::InvalidOidc)));
}
#[tokio::test]
async fn names_excluded_by_hosted_keys_read_the_process_environment() {
let name = "LITELLM_RUST_BRIDGE_UNDECLARED";
unsafe { std::env::set_var(name, "env-key") };
let server = MockServer::start().await;
let result = resolve(
state(
&server,
KeyManagementSettings {
hosted_keys: Some(vec!["OTHER".into()]),
..Default::default()
},
),
name,
)
.await;
unsafe { std::env::remove_var(name) };
assert_eq!(result.as_deref(), Some("env-key"));
assert_eq!(server.received_requests().await.unwrap().len(), 0);
}
}

View file

@ -7,14 +7,11 @@ use litellm_secrets::{
read_secret_from_python_manager,
};
use litellm_secrets_types::PythonSecretRead;
use pyo3::{exceptions::PyAttributeError, prelude::*, types::PyDict};
use pyo3::{prelude::*, types::PyDict};
use super::ManagerError;
use crate::execution::{run_async_value, run_sync_value};
const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager";
#[derive(Clone, PartialEq)]
struct Configuration {
system: KeyManagementSystem,
settings: KeyManagementSettings,
@ -90,100 +87,6 @@ impl NativeSecretManager {
)
}
#[staticmethod]
pub(super) fn from_client(client: &Bound<'_, PyAny>) -> PyResult<Option<Py<Self>>> {
let py = client.py();
if let Ok(native) = client.extract::<Py<Self>>() {
native.borrow(py).backend()?;
return Ok(Some(native));
}
let config = py
.import(SECRET_MANAGER_MODULE)?
.getattr("native_secret_manager_config")?
.call1((client,))?;
if config.is_none() {
return Ok(None);
}
if !config.getattr("owner_type")?.is(client.get_type()) {
return Ok(None);
}
let methods = config
.getattr("methods")?
.extract::<Vec<(String, Py<PyAny>)>>()?;
for (name, original) in methods {
let current = client.getattr(name.as_str())?;
let implementation = optional_attribute(&current, "__func__")?.unwrap_or(current);
if !implementation.is(original.bind(py)) {
return Ok(None);
}
}
let environment_attributes: BTreeMap<String, String> = config
.getattr("environment_attributes")?
.extract::<Vec<(String, String)>>()?
.into_iter()
.collect();
let captured = config
.getattr("environment")?
.extract::<Vec<(String, String)>>()?;
let overrides = environment_attributes
.iter()
.map(|(key, attribute)| {
let value = attribute_path(client, attribute)?;
Ok((
key.clone(),
if value.is_none() {
None
} else {
Some(value.str()?.extract::<String>()?)
},
))
})
.collect::<PyResult<Vec<_>>>()?;
let settings =
from_py::<serde_json::Map<String, serde_json::Value>>(&config.getattr("settings")?)
.map_err(ManagerError::settings)?;
let attributes = config
.getattr("settings_attributes")?
.extract::<Vec<String>>()?;
let setting_overrides = attributes
.into_iter()
.map(|name| {
let value = from_py::<serde_json::Value>(&client.getattr(name.as_str())?)?;
Ok((name, value))
})
.collect::<PyResult<Vec<_>>>()?;
let configuration = Configuration {
system: serde_json::from_value(serde_json::Value::String(
config.getattr("system")?.extract()?,
))
.map_err(ManagerError::UnknownSystem)?,
settings: serde_json::from_value(serde_json::Value::Object(
settings.into_iter().chain(setting_overrides).collect(),
))
.map_err(ManagerError::settings)?,
environment: captured
.into_iter()
.filter(|(key, _)| !environment_attributes.contains_key(key))
.chain(
overrides
.into_iter()
.filter_map(|(key, value)| value.map(|value| (key, value))),
)
.collect(),
enterprise_enabled: config.getattr("enterprise_enabled")?.extract()?,
};
if let Some(native) = cached(client, &configuration)? {
return Ok(Some(native));
}
let runtime = Self::build(py, configuration)?;
if let Some(native) = cached(client, &runtime.configuration)? {
return Ok(Some(native));
}
let native = Py::new(py, runtime)?;
client.setattr("_litellm_native_secret_manager", native.bind(py))?;
Ok(Some(native))
}
#[getter]
fn system(&self) -> &'static str {
self.configuration.system.into()
@ -395,37 +298,6 @@ impl NativeSecretManager {
}
}
fn optional_attribute<'py>(
object: &Bound<'py, PyAny>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match object.getattr(name) {
Ok(value) => Ok(Some(value)),
Err(error) if error.is_instance_of::<PyAttributeError>(object.py()) => Ok(None),
Err(error) => Err(error),
}
}
fn attribute_path<'py>(object: &Bound<'py, PyAny>, path: &str) -> PyResult<Bound<'py, PyAny>> {
match path.split_once('.') {
Some((head, tail)) => attribute_path(&object.getattr(head)?, tail),
None => object.getattr(path),
}
}
fn cached(
client: &Bound<'_, PyAny>,
configuration: &Configuration,
) -> PyResult<Option<Py<NativeSecretManager>>> {
let Some(value) = optional_attribute(client, "_litellm_native_secret_manager")? else {
return Ok(None);
};
let native = value.extract::<Py<NativeSecretManager>>()?;
let same_configuration = native.borrow(client.py()).pid == std::process::id()
&& &native.borrow(client.py()).configuration == configuration;
Ok(same_configuration.then_some(native))
}
fn parse_settings(value: Option<&Bound<'_, PyAny>>) -> PyResult<KeyManagementSettings> {
value
.map(|value| {

View file

@ -7758,7 +7758,6 @@ class ProxyConfig:
"""
Initialize the relevant secret manager if `key_management_system` is provided
"""
previous_client: Final[object] = litellm.secret_manager_client
if key_management_system is not None:
if key_management_system == KeyManagementSystem.AZURE_KEY_VAULT.value:
### LOAD FROM AZURE KEY VAULT ###
@ -7807,11 +7806,6 @@ class ProxyConfig:
else:
raise ValueError("Invalid Key Management System selected")
from litellm.rust_bridge.host.secret_manager import capture_secret_manager
if litellm.secret_manager_client is not previous_client:
capture_secret_manager(litellm.secret_manager_client, key_management_system)
def get_model_info_with_id(self, model, db_model=False) -> RouterModelInfo:
"""
Common logic across add + delete router models

View file

@ -252,8 +252,6 @@ class _SecretManagerRuntime:
settings: Mapping[str, object] | None = None,
enterprise_enabled: bool = False,
) -> _SecretManagerRuntime: ...
@staticmethod
def from_client(client: object) -> _SecretManagerRuntime | None: ...
@property
def system(self) -> str: ...
def read_secret(self, name: str, settings: Mapping[str, object] | None = None) -> JsonValue: ...

View file

@ -2,7 +2,7 @@
Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds)
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (a batch of secrets read through `get_secret_str`), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds)
- A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here
- Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers
- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green

View file

@ -1,165 +1,7 @@
from __future__ import annotations
import os
from collections.abc import Mapping
from dataclasses import dataclass, field
from importlib import import_module
from typing import Final
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
@dataclass(frozen=True, slots=True)
class NativeSecretManagerConfig:
system: str
environment: tuple[tuple[str, str], ...] = field(repr=False)
settings: Mapping[str, object] = field(repr=False)
enterprise_enabled: bool
owner_type: type[object]
environment_attributes: tuple[tuple[str, str], ...]
settings_attributes: tuple[str, ...]
methods: tuple[tuple[str, object], ...] = field(repr=False)
@dataclass(frozen=True, slots=True)
class _ClientAdapter:
system: KeyManagementSystem
module: str
name: str
methods: tuple[str, ...]
environment_attributes: tuple[tuple[str, str], ...] = ()
settings_attributes: tuple[str, ...] = ()
enterprise_enabled: bool = False
_ADAPTERS: Final = (
_ClientAdapter(
KeyManagementSystem.AWS_SECRET_MANAGER,
"litellm.secret_managers.aws_secret_manager_v2",
"AWSSecretsManagerV2",
("sync_read_secret", "async_read_secret"),
settings_attributes=(
"aws_region_name",
"aws_role_name",
"aws_session_name",
"aws_external_id",
"aws_profile_name",
"aws_web_identity_token",
"aws_sts_endpoint",
"replica_regions",
"kms_key_id",
),
),
_ClientAdapter(
KeyManagementSystem.HASHICORP_VAULT,
"litellm.secret_managers.hashicorp_secret_manager",
"HashicorpSecretManager",
("sync_read_secret", "async_read_secret", "async_write_secret", "async_delete_secret", "async_rotate_secret"),
environment_attributes=(
("HCP_VAULT_ADDR", "vault_addr"),
("HCP_VAULT_TOKEN", "vault_token"),
("HCP_VAULT_NAMESPACE", "vault_namespace"),
("HCP_VAULT_LOGIN_NAMESPACE", "login_namespace_override"),
("HCP_VAULT_SECRET_NAMESPACE", "secret_namespace_override"),
("HCP_VAULT_MOUNT_NAME", "vault_mount_name"),
("HCP_VAULT_PATH_PREFIX", "vault_path_prefix"),
("HCP_VAULT_CLIENT_CERT", "tls_cert_path"),
("HCP_VAULT_CLIENT_KEY", "tls_key_path"),
("HCP_VAULT_CERT_ROLE", "vault_cert_role"),
("HCP_VAULT_APPROLE_ROLE_ID", "approle_role_id"),
("HCP_VAULT_APPROLE_SECRET_ID", "approle_secret_id"),
("HCP_VAULT_APPROLE_MOUNT_PATH", "approle_mount_path"),
("HCP_VAULT_REFRESH_INTERVAL", "cache.default_ttl"),
),
enterprise_enabled=True,
),
_ClientAdapter(
KeyManagementSystem.CYBERARK,
"litellm.secret_managers.cyberark_secret_manager",
"CyberArkSecretManager",
("sync_read_secret", "async_read_secret", "async_write_secret", "async_delete_secret", "async_rotate_secret"),
environment_attributes=(
("CYBERARK_API_BASE", "conjur_addr"),
("CYBERARK_ACCOUNT", "conjur_account"),
("CYBERARK_USERNAME", "conjur_username"),
("CYBERARK_API_KEY", "conjur_api_key"),
("CYBERARK_CLIENT_CERT", "tls_cert_path"),
("CYBERARK_CLIENT_KEY", "tls_key_path"),
("CYBERARK_SSL_VERIFY", "ssl_verify"),
("CYBERARK_REFRESH_INTERVAL", "cache.default_ttl"),
),
enterprise_enabled=True,
),
_ClientAdapter(
KeyManagementSystem.GOOGLE_SECRET_MANAGER,
"litellm.secret_managers.google_secret_manager",
"GoogleSecretManager",
("get_secret_from_google_secret_manager",),
environment_attributes=(
("GOOGLE_SECRET_MANAGER_PROJECT_ID", "PROJECT_ID"),
("GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL", "cache.default_ttl"),
("GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", "always_read_secret_manager"),
),
enterprise_enabled=True,
),
)
_SDK_ADAPTERS: Final = (
_ClientAdapter(KeyManagementSystem.AZURE_KEY_VAULT, "azure.keyvault.secrets", "SecretClient", ("get_secret",)),
_ClientAdapter(KeyManagementSystem.GOOGLE_KMS, "google.cloud.kms_v1", "KeyManagementServiceClient", ("decrypt",)),
)
def _capture(client: object, adapter: _ClientAdapter) -> NativeSecretManagerConfig:
prefixes: Final = ("AWS_", "AZURE_", "GOOGLE_", "VERTEX_", "GCS_", "HCP_VAULT_", "CYBERARK_")
config: Final = NativeSecretManagerConfig(
system=adapter.system.value,
environment=tuple(
(name, value)
for name, value in os.environ.items()
if name.startswith(prefixes) or name == "SECRET_MANAGER_REFRESH_INTERVAL"
),
settings=KeyManagementSettings().model_dump(mode="json"),
enterprise_enabled=adapter.enterprise_enabled,
owner_type=type(client),
environment_attributes=adapter.environment_attributes,
settings_attributes=adapter.settings_attributes,
methods=tuple((name, getattr(type(client), name)) for name in adapter.methods),
)
vars(client)["_litellm_native_secret_config"] = config
return config
def capture_secret_manager(client: object, system: str) -> None:
for adapter in (*_ADAPTERS, *_SDK_ADAPTERS):
if (
adapter.system.value == system
and type(client).__name__ == adapter.name
and type(client) is getattr(import_module(adapter.module), adapter.name)
):
_capture(client, adapter)
return
if (
system == KeyManagementSystem.AWS_KMS.value
and type(client).__module__ == "botocore.client"
and type(client).__name__ == "KMS"
):
_capture(client, _ClientAdapter(KeyManagementSystem.AWS_KMS, "botocore.client", "KMS", ("decrypt",)))
def native_secret_manager_config(client: object) -> NativeSecretManagerConfig | None:
captured: Final = getattr(client, "_litellm_native_secret_config", None)
if isinstance(captured, NativeSecretManagerConfig):
return captured
for adapter in _ADAPTERS:
if type(client).__module__ == adapter.module and type(client) is getattr(
import_module(adapter.module), adapter.name
):
return _capture(client, adapter)
return None
def read_secret_str(name: str) -> str | None:
def read_secrets(names: tuple[str, ...]) -> tuple[str | None, ...]:
from litellm.secret_managers.main import get_secret_str
return get_secret_str(name)
return tuple(get_secret_str(name) for name in names)

View file

@ -1,7 +1,6 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Final
@dataclass(frozen=True, slots=True)
@ -31,83 +30,6 @@ class ProviderDefaults:
enable_azure_ad_token_refresh: object
@dataclass(frozen=True, slots=True)
class SecretManager:
readable: bool
native: bool
@dataclass(frozen=True, slots=True)
class SecretManagerBinding:
system: object
access_mode: object
hosted_keys: object
primary_secret_name: object
store_virtual_keys: object
prefix_for_stored_virtual_keys: object
kms_key_id: object
custom_secret_manager: object
aws_region_name: object
aws_role_name: object
aws_session_name: object
aws_external_id: object
aws_profile_name: object
aws_web_identity_token: object
aws_sts_endpoint: object
replica_regions: object
client: object
settings_object: object
def secret_manager() -> SecretManager:
import litellm
from litellm.secret_managers.main import (
_should_read_secret_from_secret_manager, # pyright: ignore[reportPrivateUsage] # canonical resolver is private
)
readable: Final = _should_read_secret_from_secret_manager()
system: Final = (
litellm._key_management_system # pyright: ignore[reportPrivateUsage] # canonical key management globals are private
)
return SecretManager(readable=readable, native=readable and system is not None)
def secret_manager_binding() -> SecretManagerBinding:
import litellm
from litellm.types.secret_managers.main import KeyManagementSettings
configured_system: Final = (
litellm._key_management_system # pyright: ignore[reportPrivateUsage] # canonical key management globals are private
)
configured_settings: Final = (
litellm._key_management_settings # pyright: ignore[reportPrivateUsage] # canonical key management globals are private
)
settings: Final = configured_settings or KeyManagementSettings()
system: Final = (
configured_system.value if litellm.secret_manager_client is not None and configured_system is not None else None
)
return SecretManagerBinding(
system=system,
access_mode=settings.access_mode,
hosted_keys=settings.hosted_keys,
primary_secret_name=settings.primary_secret_name,
store_virtual_keys=settings.store_virtual_keys,
prefix_for_stored_virtual_keys=settings.prefix_for_stored_virtual_keys,
kms_key_id=settings.kms_key_id,
custom_secret_manager=settings.custom_secret_manager,
aws_region_name=settings.aws_region_name,
aws_role_name=settings.aws_role_name,
aws_session_name=settings.aws_session_name,
aws_external_id=settings.aws_external_id,
aws_profile_name=settings.aws_profile_name,
aws_web_identity_token=settings.aws_web_identity_token,
aws_sts_endpoint=settings.aws_sts_endpoint,
replica_regions=settings.replica_regions,
client=litellm.secret_manager_client,
settings_object=configured_settings,
)
def provider_defaults() -> ProviderDefaults:
import litellm

View file

@ -2553,7 +2553,7 @@ def test_ProxyConfig__load_environment_variables_blocks_dangerous_keys(monkeypat
@pytest.mark.parametrize(
("flag", "system"), (("use_google_kms", "google_kms"), ("use_azure_key_vault", "azure_key_vault"))
)
async def test_load_config_legacy_secret_manager_flags_capture_the_initialized_client(
async def test_load_config_legacy_secret_manager_flags_initialize_the_client(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, flag: str, system: str
) -> None:
if system == "azure_key_vault":
@ -2561,8 +2561,6 @@ async def test_load_config_legacy_secret_manager_flags_capture_the_initialized_c
else:
client_type: Final = pytest.importorskip("google.cloud.kms_v1").KeyManagementServiceClient
from litellm.rust_bridge.host.secret_manager import native_secret_manager_config
credentials_file: Final = tmp_path / "credentials.json"
credentials_file.write_text(
json.dumps(
@ -2594,10 +2592,6 @@ async def test_load_config_legacy_secret_manager_flags_capture_the_initialized_c
client: Final = litellm.secret_manager_client
assert isinstance(client, client_type)
try:
captured: Final = native_secret_manager_config(client)
assert captured is not None
assert captured.system == system
assert dict(captured.environment)["GOOGLE_APPLICATION_CREDENTIALS"] == str(credentials_file)
assert litellm._key_management_system is not None
assert litellm._key_management_system.value == system
finally:

View file

@ -2,7 +2,7 @@
Test what each side of the bridge does, not the rollout policy that picks a side. `LITELLM_RUST` and `catalog.RULES` change every time a route or backend rolls forward, so a test that sets the env var or patches the catalog to reach a path goes red on a policy change even when the code under test is fine
Call each path directly with an explicit decision instead. The Python path is the implementation the dispatcher falls back to, e.g. `litellm.responses.main.responses`. The Rust path is the native binding, e.g. `NATIVE_OCR.load()` from `litellm/rust_bridge/ocr/entrypoints.py`, called with the `NativeCall` envelope that dispatch would hand it for every public inference route. When the native side reads a policy-derived setting such as `settings.secret_manager().native`, pin that field in the test instead of deriving it from the catalog. `ocr/test_secrets.py` shows the pattern
Call each path directly with an explicit decision instead. The Python path is the implementation the dispatcher falls back to, e.g. `litellm.responses.main.responses`. The Rust path is the native binding, e.g. `NATIVE_OCR.load()` from `litellm/rust_bridge/ocr/entrypoints.py`, called with the `NativeCall` envelope that dispatch would hand it for every public inference route. When the native side reads a policy-derived setting, pin that field in the test instead of deriving it from the catalog
Rollout policy itself, meaning which rule matches and what `LITELLM_RUST` changes, belongs in `test_catalog.py` and `test_configuration.py`, tested against rules the test builds rather than the shipped `catalog.RULES`. Which path a public call takes under those rules is a dispatch test and lives in `tests/test_litellm_rust/` for now

View file

@ -3,9 +3,10 @@ from __future__ import annotations
import asyncio
import inspect
import json
import os
import re
from functools import partial
from types import SimpleNamespace
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from urllib.parse import urlsplit
@ -18,7 +19,6 @@ from pydantic import JsonValue
import litellm
from litellm.rust_bridge import _native as native
from litellm.rust_bridge.host.secret_manager import capture_secret_manager, native_secret_manager_config
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager
from litellm.secret_managers.hashicorp_secret_manager import HashicorpSecretManager
@ -67,10 +67,24 @@ def _vault_body(value: str) -> dict[str, object]:
}
def _native_handle(client: object) -> "_SecretManagerRuntime":
handle: Final = native._SecretManagerRuntime.from_client(client)
assert handle is not None
return handle
_NATIVE_SYSTEMS: Final = MappingProxyType(
{
AWSSecretsManagerV2: KeyManagementSystem.AWS_SECRET_MANAGER,
HashicorpSecretManager: KeyManagementSystem.HASHICORP_VAULT,
CyberArkSecretManager: KeyManagementSystem.CYBERARK,
}
)
def _native_handle(
client: AWSSecretsManagerV2 | HashicorpSecretManager | CyberArkSecretManager,
) -> _SecretManagerRuntime:
return native._SecretManagerRuntime.from_config(
_NATIVE_SYSTEMS[type(client)].value,
{name: value for name, value in os.environ.items() if name.startswith(("AWS_", "HCP_VAULT_", "CYBERARK_"))},
KeyManagementSettings(aws_region_name="us-east-1").model_dump(mode="json"),
enterprise_enabled=True,
)
@pytest.mark.parametrize("system", ("aws_secret_manager", "hashicorp_vault", "cyberark"))
@ -107,64 +121,6 @@ async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(s
assert handle.read_secret("KEY") == expected
def test_shared_initializer_captures_credentials_and_tracks_instance_settings(monkeypatch: pytest.MonkeyPatch) -> None:
with recording_service() as server:
server.default_response = ResponseSpec(body={"SecretString": "native-value"})
server.expected_requests = 2
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "captured-access")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "captured-secret")
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", server.base_url)
from litellm.proxy.proxy_server import ProxyConfig
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(aws_region_name="us-east-1"))
ProxyConfig().initialize_secret_manager(KeyManagementSystem.AWS_SECRET_MANAGER.value)
manager: Final = litellm.secret_manager_client
assert isinstance(manager, AWSSecretsManagerV2)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "changed-access")
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", "http://127.0.0.1:1")
first: Final = native._SecretManagerRuntime.from_client(manager)
assert first is not None
assert native._SecretManagerRuntime.from_client(manager) is first
assert first.read_secret("KEY") == "native-value"
manager.aws_region_name = "us-west-2"
second: Final = native._SecretManagerRuntime.from_client(manager)
assert second is not None
assert second is not first
assert second.read_secret("KEY") == "native-value"
assert "Credential=captured-access/" in server.requests[0].headers["authorization"]
assert "/us-east-1/" in server.requests[0].headers["authorization"]
assert "/us-west-2/" in server.requests[1].headers["authorization"]
def test_configuration_replacement_rebuilds_without_invalidating_existing_handles(
monkeypatch: pytest.MonkeyPatch,
) -> None:
with recording_service() as first_server, recording_service() as second_server:
first_server.default_response = ResponseSpec(body=_vault_body("first"))
second_server.default_response = ResponseSpec(body=_vault_body("second"))
manager: Final = _vault(monkeypatch, first_server.base_url)
first: Final = native._SecretManagerRuntime.from_client(manager)
assert first is not None
assert first.read_secret("KEY") == "first"
manager.vault_addr = second_server.base_url
second: Final = native._SecretManagerRuntime.from_client(manager)
assert second is not None
assert second is not first
assert second.read_secret("KEY") == "second"
assert first.read_secret("KEY") == "first"
def test_custom_subclass_is_not_captured_for_a_native_handle() -> None:
class CustomManager(AWSSecretsManagerV2):
def sync_read_secret(self, secret_name: str, primary_secret_name: str | None = None) -> str:
return f"custom:{secret_name}"
manager: Final = CustomManager(aws_region_name="us-east-1")
assert native._SecretManagerRuntime.from_client(manager) is None
assert manager.sync_read_secret("KEY") == "custom:KEY"
@pytest.mark.parametrize("system", ("google_secret_manager", "hashicorp_vault", "cyberark"))
def test_enterprise_backends_cannot_initialize_without_entitlement(system: str) -> None:
with pytest.raises(ValueError, match=r"[Ee]nterprise|[Pp]remium"):
@ -176,19 +132,6 @@ def test_azure_factory_rejects_unencrypted_vault_endpoints() -> None:
native._SecretManagerRuntime.from_config("azure_key_vault", {"AZURE_KEY_VAULT_URI": "http://127.0.0.1:1"})
def test_direct_builtin_constructor_can_use_native_without_registration(monkeypatch: pytest.MonkeyPatch) -> None:
with recording_service() as server:
server.default_response = ResponseSpec(body=_vault_body("native-value"))
manager: Final = _vault(monkeypatch, server.base_url)
handle: Final = native._SecretManagerRuntime.from_client(manager)
assert handle is not None
assert handle.read_secret("KEY") == "native-value"
def test_uncaptured_clients_have_no_native_handle() -> None:
assert native._SecretManagerRuntime.from_client(SimpleNamespace()) is None
def test_native_failure_is_not_replayed_in_python(monkeypatch: pytest.MonkeyPatch) -> None:
with recording_service() as server:
server.default_response = ResponseSpec(status=403, body={"errors": ["denied"]})
@ -200,62 +143,6 @@ def test_native_failure_is_not_replayed_in_python(monkeypatch: pytest.MonkeyPatc
assert len(server.requests) == 1
@pytest.mark.parametrize("explicit_capture", (False, True))
def test_config_capture_preserves_credentials_and_excludes_unrelated_environment(
monkeypatch: pytest.MonkeyPatch, explicit_capture: bool
) -> None:
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "initial-access")
monkeypatch.setenv("SECRET_MANAGER_REFRESH_INTERVAL", "45")
monkeypatch.setenv("UNRELATED_PRIVATE_TOKEN", "unrelated-secret")
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
if explicit_capture:
capture_secret_manager(manager, "aws_secret_manager")
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "replacement-access")
captured: Final = native_secret_manager_config(manager)
assert captured is not None
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "replacement-access")
retained: Final = native_secret_manager_config(manager)
assert retained is captured
assert retained.system == "aws_secret_manager"
assert dict(retained.environment)["AWS_ACCESS_KEY_ID"] == "initial-access"
assert dict(retained.environment)["SECRET_MANAGER_REFRESH_INTERVAL"] == "45"
assert "UNRELATED_PRIVATE_TOKEN" not in dict(retained.environment)
assert "initial-access" not in repr(retained)
assert retained.settings == KeyManagementSettings().model_dump(mode="json")
def test_kms_sdk_client_capture_preserves_environment(monkeypatch: pytest.MonkeyPatch) -> None:
import boto3
monkeypatch.setenv("AWS_REGION_NAME", "captured-region")
client: Final = boto3.client(
"kms", region_name="us-east-1", aws_access_key_id="access", aws_secret_access_key="secret"
)
try:
capture_secret_manager(client, "aws_kms")
monkeypatch.setenv("AWS_REGION_NAME", "replacement-region")
captured: Final = native_secret_manager_config(client)
assert captured is not None
assert captured.system == "aws_kms"
assert dict(captured.environment)["AWS_REGION_NAME"] == "captured-region"
finally:
client.close()
def test_same_named_custom_client_is_not_captured() -> None:
class AWSSecretsManagerV2:
pass
client: Final = AWSSecretsManagerV2()
capture_secret_manager(client, "aws_secret_manager")
assert native_secret_manager_config(client) is None
@pytest.mark.parametrize("value", ("text", "", True, False, 42, 2**100, [1, "two"], {"nested": True}, None))
async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.MonkeyPatch, value: JsonValue) -> None:
with recording_service() as server:
@ -267,8 +154,7 @@ async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.Monke
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
settings: Final = KeyManagementSettings(primary_secret_name="primary")
reference: Final = get_secret_from_manager(manager, "aws_secret_manager", "KEY", settings)
handle: Final = native._SecretManagerRuntime.from_client(manager)
assert handle is not None
handle: Final = _native_handle(manager)
actual: Final = handle.read_secret("KEY", settings.model_dump(mode="json"))
asynchronous: Final = await handle.read_secret_async("KEY", settings.model_dump(mode="json"))
assert type(actual) is type(reference) is type(value)
@ -295,7 +181,9 @@ async def test_aws_primary_values_match_python_reads(monkeypatch: pytest.MonkeyP
asynchronous: Final = await handle.read_secret_async("KEY", settings_dump)
assert synchronous == reference
assert asynchronous == reference
assert tuple(json.loads(request.raw_body) for request in server.requests) == ({"SecretId": primary or "KEY"},) * 3
assert (
tuple(json.loads(request.raw_body) for request in server.requests) == ({"SecretId": primary or "KEY"},) * 3
)
@pytest.mark.parametrize("primary", (None, "primary"))

View file

@ -1,14 +1,10 @@
from typing import Final
import httpx
import pytest
import litellm
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.llms.custom_httpx.http_handler import default_user_agent
from litellm.rust_bridge.host import settings
from litellm.secret_managers.main import get_secret_str
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
def test_url_policy_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None:
@ -57,143 +53,6 @@ def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyP
assert result.ssl_verify is True
class _VaultSecrets(CustomSecretManager):
def __init__(self, secrets: dict[str, str]) -> None:
super().__init__(secret_manager_name="rust_bridge_settings_test")
self.secrets = secrets
async def async_read_secret(
self,
secret_name: str,
optional_params: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
) -> str | None:
return self.secrets.get(secret_name)
def sync_read_secret(
self,
secret_name: str,
optional_params: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
) -> str | None:
return self.secrets.get(secret_name)
@pytest.mark.parametrize(
("access_mode", "readable"),
[("read_only", True), ("read_and_write", True), ("write_only", False)],
)
def test_secret_manager_is_readable_only_when_litellm_would_read_secrets_from_it(
monkeypatch: pytest.MonkeyPatch, access_mode: str, readable: bool
) -> None:
monkeypatch.setenv("MISTRAL_API_KEY", "env-key")
monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets({"MISTRAL_API_KEY": "vault-key"}))
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode))
assert settings.secret_manager() == settings.SecretManager(readable=readable, native=readable)
assert (get_secret_str("MISTRAL_API_KEY") == "vault-key") is readable
@pytest.mark.parametrize(
("system", "access_mode", "native"),
[
(KeyManagementSystem.CUSTOM, "read_only", True),
(KeyManagementSystem.AWS_SECRET_MANAGER, "read_only", True),
(KeyManagementSystem.CUSTOM, "write_only", False),
(None, "read_only", False),
],
)
def test_secret_manager_is_native_only_when_a_readable_system_is_configured(
monkeypatch: pytest.MonkeyPatch,
system: KeyManagementSystem | None,
access_mode: str,
native: bool,
) -> None:
monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets({}))
monkeypatch.setattr(litellm, "_key_management_system", system)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode))
assert settings.secret_manager() == settings.SecretManager(
readable=access_mode != "write_only", native=native
)
def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "secret_manager_client", None)
assert settings.secret_manager() == settings.SecretManager(readable=False, native=False)
def test_secret_manager_projects_custom_settings(monkeypatch: pytest.MonkeyPatch) -> None:
manager_settings: Final = KeyManagementSettings(
access_mode="read_and_write",
hosted_keys=["MISTRAL_API_KEY"],
primary_secret_name="primary",
aws_region_name="us-east-1",
)
client: Final = _VaultSecrets({"MISTRAL_API_KEY": "vault-key"})
monkeypatch.setattr(litellm, "secret_manager_client", client)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
monkeypatch.setattr(litellm, "_key_management_settings", manager_settings)
assert settings.secret_manager_binding() == settings.SecretManagerBinding(
system="custom",
access_mode="read_and_write",
hosted_keys=["MISTRAL_API_KEY"],
primary_secret_name="primary",
store_virtual_keys=manager_settings.store_virtual_keys,
prefix_for_stored_virtual_keys=manager_settings.prefix_for_stored_virtual_keys,
kms_key_id=manager_settings.kms_key_id,
custom_secret_manager=manager_settings.custom_secret_manager,
aws_region_name="us-east-1",
aws_role_name=manager_settings.aws_role_name,
aws_session_name=manager_settings.aws_session_name,
aws_external_id=manager_settings.aws_external_id,
aws_profile_name=manager_settings.aws_profile_name,
aws_web_identity_token=manager_settings.aws_web_identity_token,
aws_sts_endpoint=manager_settings.aws_sts_endpoint,
replica_regions=manager_settings.replica_regions,
client=client,
settings_object=manager_settings,
)
def test_secret_manager_without_a_client_has_no_system(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "secret_manager_client", None)
assert settings.secret_manager_binding().system is None
def test_secret_manager_uses_key_management_defaults(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "secret_manager_client", None)
monkeypatch.setattr(litellm, "_key_management_settings", None)
defaults: Final = KeyManagementSettings()
result: Final = settings.secret_manager_binding()
assert result == settings.SecretManagerBinding(
system=None,
access_mode=defaults.access_mode,
hosted_keys=defaults.hosted_keys,
primary_secret_name=defaults.primary_secret_name,
store_virtual_keys=defaults.store_virtual_keys,
prefix_for_stored_virtual_keys=defaults.prefix_for_stored_virtual_keys,
kms_key_id=defaults.kms_key_id,
custom_secret_manager=defaults.custom_secret_manager,
aws_region_name=defaults.aws_region_name,
aws_role_name=defaults.aws_role_name,
aws_session_name=defaults.aws_session_name,
aws_external_id=defaults.aws_external_id,
aws_profile_name=defaults.aws_profile_name,
aws_web_identity_token=defaults.aws_web_identity_token,
aws_sts_endpoint=defaults.aws_sts_endpoint,
replica_regions=defaults.replica_regions,
client=None,
settings_object=None,
)
def test_provider_defaults_read_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "vertex_project", "configured-project")
monkeypatch.setattr(litellm, "vertex_location", "europe-west4")

View file

@ -1,7 +1,6 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping
from dataclasses import replace
from typing import Final, Protocol, cast # noqa: TID251 # narrows the parametrized path to its protocol
import httpx
@ -10,7 +9,6 @@ import pytest
import litellm
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.llms.anthropic.messages.native.handler import anthropic_messages
from litellm.rust_bridge.host import settings
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES
from litellm.rust_bridge.public_call import NativeCall
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
@ -97,8 +95,6 @@ async def test_secret_manager_supplies_the_anthropic_key_and_base(
)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
configured: Final = settings.secret_manager
monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=True))
await messages()

View file

@ -3,7 +3,6 @@ from __future__ import annotations
import asyncio
from collections.abc import Awaitable, Generator, Mapping
from contextlib import contextmanager
from dataclasses import replace
from typing import Final, Literal, Protocol, TypeAlias, cast
import httpx
@ -12,8 +11,6 @@ import pytest
import litellm
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge import _native as native
from litellm.rust_bridge.host import settings
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
from litellm.rust_bridge.public_call import NativeCall
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
@ -119,7 +116,6 @@ def _configure(
*,
manager: _VaultSecrets,
key_management: KeyManagementSettings,
native_secret_manager: bool = True,
environment_key: str | None = "environment-key",
) -> None:
if environment_key is None:
@ -129,8 +125,6 @@ def _configure(
monkeypatch.setattr(litellm, "secret_manager_client", manager)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
monkeypatch.setattr(litellm, "_key_management_settings", key_management)
configured: Final = settings.secret_manager
monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=native_secret_manager))
@pytest.mark.parametrize(
@ -220,7 +214,7 @@ async def test_rust_reads_a_python_only_secret_manager_through_python(
) -> None:
manager: Final = _VaultSecrets()
key_management: Final = KeyManagementSettings(access_mode="read_only")
_configure(monkeypatch, manager=manager, key_management=key_management, native_secret_manager=False)
_configure(monkeypatch, manager=manager, key_management=key_management)
with _mistral_service() as server:
await rust_ocr(server.base_url)
@ -245,7 +239,7 @@ class _PythonManagerReads:
return "python-key" if secret_name == "MISTRAL_API_KEY" else None
async def test_python_only_builtin_manager_is_read_through_python_not_its_native_backend(
async def test_builtin_manager_is_read_through_python_not_its_native_backend(
monkeypatch: pytest.MonkeyPatch, rust_ocr: Ocr
) -> None:
from litellm.secret_managers import main as secret_manager_main
@ -255,19 +249,18 @@ async def test_python_only_builtin_manager_is_read_through_python_not_its_native
monkeypatch.setattr(secret_manager_main, "get_secret_from_manager", python_reads)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "native-access")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "native-secret")
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", "http://127.0.0.1:1")
monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
monkeypatch.setattr(litellm, "secret_manager_client", manager)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.AWS_SECRET_MANAGER)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(hosted_keys=["MISTRAL_API_KEY"]))
monkeypatch.setattr(settings, "secret_manager", lambda: settings.SecretManager(readable=True, native=False))
with _mistral_service() as provider:
await rust_ocr(provider.base_url)
assert provider.requests[0].headers["authorization"] == "Bearer python-key"
assert (manager, "aws_secret_manager", "MISTRAL_API_KEY") in python_reads.reads
assert getattr(manager, "_litellm_native_secret_manager", None) is None
async def test_no_secret_client_leaves_dormant_binding_settings_unread(
@ -341,8 +334,6 @@ def _configure_client(
monkeypatch.setattr(litellm, "secret_manager_client", client)
monkeypatch.setattr(litellm, "_key_management_system", system)
monkeypatch.setattr(litellm, "_key_management_settings", key_management)
configured: Final = settings.secret_manager
monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=True))
async def _assert_missing_key(ocr: Ocr) -> None:
@ -401,7 +392,9 @@ async def test_manager_without_the_key_does_not_fall_back_to_the_environment(
await _assert_missing_key(ocr)
async def test_rust_hosted_keys_exclude_azure_sdk_clients_too(monkeypatch: pytest.MonkeyPatch, rust_ocr: Ocr) -> None:
async def test_azure_sdk_clients_are_read_even_for_names_outside_hosted_keys(
monkeypatch: pytest.MonkeyPatch, ocr: Ocr
) -> None:
_configure_client(
monkeypatch,
client=_azure_sdk_client("vault-key"),
@ -410,11 +403,9 @@ async def test_rust_hosted_keys_exclude_azure_sdk_clients_too(monkeypatch: pytes
)
with _mistral_service() as server:
await rust_ocr(server.base_url)
await ocr(server.base_url)
assert server.requests[0].headers["authorization"] == "Bearer environment-key", (
"recorded divergence: Python's get_secret_from_manager recognizes Azure SDK clients by type and ignores hosted_keys"
)
assert server.requests[0].headers["authorization"] == "Bearer vault-key"
async def test_custom_system_with_a_foreign_client_falls_back_to_the_environment(
@ -431,41 +422,3 @@ async def test_custom_system_with_a_foreign_client_falls_back_to_the_environment
await ocr(server.base_url)
assert server.requests[0].headers["authorization"] == "Bearer environment-key"
async def test_native_backend_supplies_ocr_credentials_without_a_python_reader(
monkeypatch: pytest.MonkeyPatch, rust_ocr: Ocr
) -> None:
from litellm.secret_managers import main as secret_manager_main
from litellm.secret_managers import secret_manager_handler
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
def reject_python_read(
client: object,
key_manager: str,
secret_name: str,
key_management_settings: KeyManagementSettings | None = None,
) -> str | None:
raise AssertionError("Rust must read the native backend directly")
monkeypatch.setattr(secret_manager_handler, "get_secret_from_manager", reject_python_read)
monkeypatch.setattr(secret_manager_main, "get_secret_from_manager", reject_python_read)
with recording_service() as secrets, _mistral_service(expected_requests=2) as provider:
secrets.default_response = ResponseSpec(body={"SecretString": "native-key"})
secrets.expected_requests = None
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "native-access")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "native-secret")
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", secrets.base_url)
monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
monkeypatch.setattr(litellm, "secret_manager_client", manager)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.AWS_SECRET_MANAGER)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(hosted_keys=["MISTRAL_API_KEY"]))
monkeypatch.setattr(settings, "secret_manager", lambda: settings.SecretManager(readable=True, native=True))
await rust_ocr(provider.base_url)
await rust_ocr(provider.base_url)
assert len(secrets.requests) == 2, [(request.path, request.body) for request in secrets.requests]
assert all(request.headers["authorization"] == "Bearer native-key" for request in provider.requests)
assert all("Credential=native-access/" in request.headers["authorization"] for request in secrets.requests)
assert native._SecretManagerRuntime.from_client(manager) is getattr(manager, "_litellm_native_secret_manager")