From b74bd81eda5280444a0545239706d877b69ae85e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 4 Nov 2026 12:56:58 +0900 Subject: [PATCH 1/2] 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 --- litellm-rust/Cargo.lock | 2 - litellm-rust/crates/python-bridge/Cargo.toml | 2 - .../crates/python-bridge/src/coercion.rs | 28 - .../crates/python-bridge/src/errors.rs | 2 +- .../python-bridge/src/python_settings.rs | 4 - .../python-bridge/src/secrets/callback.rs | 508 ------------------ .../python-bridge/src/secrets/config.rs | 385 ------------- .../crates/python-bridge/src/secrets/error.rs | 11 +- .../crates/python-bridge/src/secrets/mod.rs | 40 +- .../python-bridge/src/secrets/python.rs | 131 +++-- .../python-bridge/src/secrets/resolved.rs | 277 ---------- .../python-bridge/src/secrets/runtime.rs | 130 +---- litellm/proxy/proxy_server.py | 6 - litellm/rust_bridge/_native.pyi | 2 - litellm/rust_bridge/host/AGENTS.md | 2 +- litellm/rust_bridge/host/secret_manager.py | 162 +----- litellm/rust_bridge/host/settings.py | 78 --- .../proxy/proxy_server/test_proxy_config.py | 8 +- tests/unit/rust_bridge/AGENTS.md | 2 +- .../rust_bridge/host/test_secret_manager.py | 160 +----- tests/unit/rust_bridge/host/test_settings.py | 141 ----- .../unit/rust_bridge/messages/test_secrets.py | 4 - tests/unit/rust_bridge/ocr/test_secrets.py | 63 +-- 23 files changed, 153 insertions(+), 1995 deletions(-) delete mode 100644 litellm-rust/crates/python-bridge/src/secrets/callback.rs delete mode 100644 litellm-rust/crates/python-bridge/src/secrets/config.rs delete mode 100644 litellm-rust/crates/python-bridge/src/secrets/resolved.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f91deaaa401..9252308c1ad 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e3ef45f7876..ad69eda2648 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/python-bridge/src/coercion.rs b/litellm-rust/crates/python-bridge/src/coercion.rs index 8cdf1e6c5bc..2ff427ca8b7 100644 --- a/litellm-rust/crates/python-bridge/src/coercion.rs +++ b/litellm-rust/crates/python-bridge/src/coercion.rs @@ -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 { - if !self.value.is_instance_of::() { - return Err(ProjectionError::InternalSchemaFailure( - self.expected("a Boolean")?, - )); - } - Ok(self.exact_true()) - } - pub(crate) fn strict_string(&self) -> Result { let Ok(value) = self.value.cast::() else { return Err(self.invalid("a string")); @@ -235,19 +220,6 @@ impl<'py> Field<'py> { }) .collect() } - - pub(crate) fn optional_string_collection( - &self, - ) -> Result>, ProjectionError> { - if self.value.is_none() { - return Ok(None); - } - self.string_collection().map(Some) - } - - pub(crate) fn python_binding(&self) -> Option> { - (!self.value.is_none()).then(|| self.value.clone().unbind()) - } } #[cfg(test)] diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 426cd861729..f3a295f03a4 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -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") }) diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index eacec6859b5..abf24bdd44b 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -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> { diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs deleted file mode 100644 index af6485a56e1..00000000000 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ /dev/null @@ -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, - context: PythonContext, -} - -struct PythonClient { - client: Py, - system: Option, - settings: Option>, -} - -impl PythonSecretManager { - pub(crate) fn new( - client: Py, - system: Option, - settings: Option>, - context: PythonContext, - ) -> Self { - Self { - client: Arc::new(PythonClient { - client, - system, - settings, - }), - context, - } - } -} - -impl PythonClient { - fn read(&self, py: Python<'_>, name: &str) -> PyResult> { - 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::() { - 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, 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::(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, 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> { - 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::().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::() - .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::().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::() - .unwrap(), - "azure_key_vault" - ); - assert_eq!( - call.get_item("secret_name") - .unwrap() - .unwrap() - .extract::() - .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, - #[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::>() - .unwrap(), - Vec::::new() - ); - let calls = locals.get_item("calls").unwrap().unwrap(); - let call = calls.get_item(0).unwrap().cast_into::().unwrap(); - assert!(call.get_item("client").unwrap().unwrap().is(&manager)); - assert_eq!( - call.get_item("key_manager") - .unwrap() - .unwrap() - .extract::() - .unwrap(), - key_manager - ); - }); - }); - } -} diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs deleted file mode 100644 index c1203dcbd89..00000000000 --- a/litellm-rust/crates/python-bridge/src/secrets/config.rs +++ /dev/null @@ -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> = - FieldSpec::new("system", parse_optional_system); -const ACCESS_MODE: FieldSpec = FieldSpec::new("access_mode", parse_access_mode); -const HOSTED_KEYS: FieldSpec>> = - FieldSpec::new("hosted_keys", |field| field.optional_string_collection()); -const STORE_VIRTUAL_KEYS: FieldSpec = - FieldSpec::new("store_virtual_keys", |field| field.truthy()); -const PREFIX_FOR_STORED_VIRTUAL_KEYS: FieldSpec = - FieldSpec::new("prefix_for_stored_virtual_keys", |field| { - field.strict_string() - }); -const PRIMARY_SECRET_NAME: FieldSpec> = - FieldSpec::new("primary_secret_name", |field| field.falsy_optional_string()); -const KMS_KEY_ID: FieldSpec> = - FieldSpec::new("kms_key_id", |field| field.falsy_optional_string()); -const CUSTOM_SECRET_MANAGER: FieldSpec> = - FieldSpec::new("custom_secret_manager", |field| { - field.falsy_optional_string() - }); -const AWS_REGION_NAME: FieldSpec> = - FieldSpec::new("aws_region_name", |field| field.falsy_optional_string()); -const AWS_ROLE_NAME: FieldSpec> = - FieldSpec::new("aws_role_name", |field| field.falsy_optional_string()); -const AWS_SESSION_NAME: FieldSpec> = - FieldSpec::new("aws_session_name", |field| field.falsy_optional_string()); -const AWS_EXTERNAL_ID: FieldSpec> = - FieldSpec::new("aws_external_id", |field| field.falsy_optional_string()); -const AWS_PROFILE_NAME: FieldSpec> = - FieldSpec::new("aws_profile_name", |field| field.falsy_optional_string()); -const AWS_WEB_IDENTITY_TOKEN: FieldSpec> = - FieldSpec::new("aws_web_identity_token", |field| { - field.falsy_optional_string() - }); -const AWS_STS_ENDPOINT: FieldSpec> = - FieldSpec::new("aws_sts_endpoint", |field| field.falsy_optional_string()); -const REPLICA_REGIONS: FieldSpec>> = - FieldSpec::new("replica_regions", |field| { - field.optional_string_collection() - }); -const CLIENT: FieldSpec>> = - FieldSpec::new("client", |field| Ok(field.python_binding())); -const SETTINGS_OBJECT: FieldSpec>> = - 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), - Native(Box), -} - -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, - /// 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>, -} - -impl SecretManagerSnapshot { - pub(crate) fn into_state(self, context: PythonContext) -> Arc { - 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 { - 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 { - 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, 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 { - 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::(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)); - }); - } -} diff --git a/litellm-rust/crates/python-bridge/src/secrets/error.rs b/litellm-rust/crates/python-bridge/src/secrets/error.rs index ca51ee79921..ebe1dcc8b64 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/error.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/error.rs @@ -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 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( diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs index 57042e6a32e..6d1677be8fd 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mod.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -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 = 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> { - 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)?)) } diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs index c5d9e3565b3..026b993fe73 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/python.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -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>, + read_secrets: Arc>, context: PythonContext, } @@ -22,18 +27,42 @@ impl PythonSecrets { pub(super) fn new(py: Python<'_>) -> PyResult { 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, context: PythonContext) -> Self { + fn reading_with(read_secrets: Py, 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) -> Result>, 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::>>()) + .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); + +impl Lookup for ReadSecrets { + fn get(&self, name: &str) -> Option { + 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, 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::>()) - .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> { + 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), #[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)) { + 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> = 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), ) { - 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)) { - secrets.0.get_secret_str("MISTRAL_API_KEY").await.unwrap(); + secrets.0.resolve(&["MISTRAL_API_KEY"]).await.unwrap(); let seen: Vec = Python::attach(|py| global(&secrets.1, py, "seen_context_values")); assert_eq!(seen, vec!["request-value".to_owned()]); diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs deleted file mode 100644 index 752b533886e..00000000000 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ /dev/null @@ -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, 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, 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 { - 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, name: &'static str) -> Option { - 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); - } -} diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index 0d52b0aa629..dc88895cc01 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -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>> { - let py = client.py(); - if let Ok(native) = client.extract::>() { - 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::)>>()?; - for (name, original) in methods { - let current = client.getattr(name.as_str())?; - let implementation = optional_attribute(¤t, "__func__")?.unwrap_or(current); - if !implementation.is(original.bind(py)) { - return Ok(None); - } - } - let environment_attributes: BTreeMap = config - .getattr("environment_attributes")? - .extract::>()? - .into_iter() - .collect(); - let captured = config - .getattr("environment")? - .extract::>()?; - 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::()?) - }, - )) - }) - .collect::>>()?; - let settings = - from_py::>(&config.getattr("settings")?) - .map_err(ManagerError::settings)?; - let attributes = config - .getattr("settings_attributes")? - .extract::>()?; - let setting_overrides = attributes - .into_iter() - .map(|name| { - let value = from_py::(&client.getattr(name.as_str())?)?; - Ok((name, value)) - }) - .collect::>>()?; - 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>> { - match object.getattr(name) { - Ok(value) => Ok(Some(value)), - Err(error) if error.is_instance_of::(object.py()) => Ok(None), - Err(error) => Err(error), - } -} - -fn attribute_path<'py>(object: &Bound<'py, PyAny>, path: &str) -> PyResult> { - 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>> { - let Some(value) = optional_attribute(client, "_litellm_native_secret_manager")? else { - return Ok(None); - }; - let native = value.extract::>()?; - 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 { value .map(|value| { diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 698c1e6e752..693c80ddd46 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8177e4a5cc9..b018aed3335 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -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: ... diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index 9b751a1ada2..7ef9c2501b6 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -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 diff --git a/litellm/rust_bridge/host/secret_manager.py b/litellm/rust_bridge/host/secret_manager.py index 4fb3350165c..44bc13227d8 100644 --- a/litellm/rust_bridge/host/secret_manager.py +++ b/litellm/rust_bridge/host/secret_manager.py @@ -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) diff --git a/litellm/rust_bridge/host/settings.py b/litellm/rust_bridge/host/settings.py index 60a58f1aaba..664af510c6f 100644 --- a/litellm/rust_bridge/host/settings.py +++ b/litellm/rust_bridge/host/settings.py @@ -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 diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 712b4134c05..403d2e44802 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -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: diff --git a/tests/unit/rust_bridge/AGENTS.md b/tests/unit/rust_bridge/AGENTS.md index abf5ac6bed7..6a85a0932fb 100644 --- a/tests/unit/rust_bridge/AGENTS.md +++ b/tests/unit/rust_bridge/AGENTS.md @@ -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 diff --git a/tests/unit/rust_bridge/host/test_secret_manager.py b/tests/unit/rust_bridge/host/test_secret_manager.py index 1f8acc5792e..4c3c5a93329 100644 --- a/tests/unit/rust_bridge/host/test_secret_manager.py +++ b/tests/unit/rust_bridge/host/test_secret_manager.py @@ -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")) diff --git a/tests/unit/rust_bridge/host/test_settings.py b/tests/unit/rust_bridge/host/test_settings.py index 5d3486d9a64..f2ac87f4557 100644 --- a/tests/unit/rust_bridge/host/test_settings.py +++ b/tests/unit/rust_bridge/host/test_settings.py @@ -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") diff --git a/tests/unit/rust_bridge/messages/test_secrets.py b/tests/unit/rust_bridge/messages/test_secrets.py index 5ee89fa7794..c7aaacd54e6 100644 --- a/tests/unit/rust_bridge/messages/test_secrets.py +++ b/tests/unit/rust_bridge/messages/test_secrets.py @@ -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() diff --git a/tests/unit/rust_bridge/ocr/test_secrets.py b/tests/unit/rust_bridge/ocr/test_secrets.py index 4ea6e09f885..421f246650f 100644 --- a/tests/unit/rust_bridge/ocr/test_secrets.py +++ b/tests/unit/rust_bridge/ocr/test_secrets.py @@ -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") From 4351906ca719dc21f7d65451e274685ba2744392 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 4 Nov 2026 13:44:06 +0900 Subject: [PATCH 2/2] style(rust-bridge): rustfmt the secrets source Co-Authored-By: Claude Opus 5.5 --- litellm-rust/crates/python-bridge/src/secrets/mod.rs | 11 +++++------ .../crates/python-bridge/src/secrets/python.rs | 12 ++++++++++-- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs index 6d1677be8fd..d1553cf5105 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mod.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -23,12 +23,11 @@ use crate::python_settings::missing_module; pub(crate) fn source(py: Python<'_>) -> PyResult> { 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) 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), } } diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs index 026b993fe73..eba1596878f 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/python.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -46,7 +46,12 @@ impl PythonSecrets { PyTuple::new(py, names) .and_then(|names| read_secrets.bind(py).call1((names,))) .and_then(|values| values.extract::>>()) - .map(|values| values.into_iter().map(|value| value.map(SecretValue::new)).collect()) + .map(|values| { + values + .into_iter() + .map(|value| value.map(SecretValue::new)) + .collect() + }) .map_err(|error| external_error(py, error)) }) .await @@ -196,7 +201,10 @@ def read_secrets(names): .await .unwrap(); - assert_eq!(resolved.get("MISTRAL_API_KEY").as_deref(), Some("vault-key")); + 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(),