fix(rust): read OCR secrets from the process environment and decline when a secret manager is readable

The OCR route called back into Python's get_secret_str for every env
fallback. With no secret manager configured that is os.environ behind a GIL
hop, and with one configured it blocked a tokio worker on vault I/O and also
sent the Azure and GCP identity variables, which Python reads with os.getenv,
to the vault. The other Rust routes already read the process environment.

Read the process environment here too. When litellm would read secrets from
a secret manager, decline the Rust route so the Python route serves the call
with the vault-backed keys
This commit is contained in:
Yujong Lee 2026-09-19 08:32:46 -07:00
parent b341d21a76
commit 1669213eb5
5 changed files with 100 additions and 86 deletions

View file

@ -19,5 +19,8 @@
"vertex_project",
"vertex_location",
"enable_azure_ad_token_refresh"
],
"secret_manager": [
"readable"
]
}

View file

@ -1,4 +1,3 @@
use litellm_core_utils::settings::Lookup;
use pyo3::prelude::*;
const MODULE: &str = "litellm.rust_bridge.settings";
@ -8,17 +7,24 @@ pub(crate) enum PythonSettings {
Http,
UrlPolicy,
ProviderDefaults,
SecretManager,
}
impl PythonSettings {
#[cfg(test)]
pub(crate) const ALL: [Self; 3] = [Self::Http, Self::UrlPolicy, Self::ProviderDefaults];
pub(crate) const ALL: [Self; 4] = [
Self::Http,
Self::UrlPolicy,
Self::ProviderDefaults,
Self::SecretManager,
];
pub(crate) fn name(self) -> &'static str {
match self {
Self::Http => "http_settings",
Self::UrlPolicy => "url_policy",
Self::ProviderDefaults => "provider_defaults",
Self::SecretManager => "secret_manager",
}
}
@ -32,23 +38,6 @@ impl PythonSettings {
}
}
pub(crate) struct PythonSecrets;
impl Lookup for PythonSecrets {
fn get(&self, name: &str) -> Option<String> {
Python::attach(|py| {
py.import(MODULE)
.and_then(|module| module.getattr("secret")?.call1((name,)))
.and_then(|value| value.extract::<Option<String>>())
.unwrap_or_else(|error| {
let _ =
PythonSettings::warn(py, &format!("reading secret {name} failed: {error}"));
None
})
})
}
}
#[cfg(test)]
pub(crate) const CONTRACT: &str = include_str!("../python_settings.json");
@ -56,10 +45,9 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json");
mod tests {
use std::{collections::BTreeSet, ffi::CString};
use litellm_core_utils::settings::Lookup;
use pyo3::{prelude::*, types::PyDict};
use super::{CONTRACT, PythonSecrets, PythonSettings};
use super::{CONTRACT, PythonSettings};
#[test]
fn every_settings_group_is_in_the_python_contract() {
@ -83,48 +71,4 @@ mod tests {
assert_eq!(read, declared);
});
}
#[test]
fn secrets_come_from_the_python_secret_reader_and_a_failed_read_is_unset() {
Python::initialize();
Python::attach(|py| {
py.run(
c"
import sys
import types
settings = types.ModuleType('litellm.rust_bridge.settings')
settings.warnings = []
def secret(name):
if name == 'BROKEN':
raise RuntimeError('vault down')
return {'MISTRAL_API_KEY': 'from-vault'}.get(name)
settings.secret = secret
settings.warn = settings.warnings.append
sys.modules.setdefault('litellm', types.ModuleType('litellm'))
sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge'))
sys.modules['litellm.rust_bridge.settings'] = settings
",
None,
None,
)
.unwrap();
});
assert_eq!(
PythonSecrets.get("MISTRAL_API_KEY").as_deref(),
Some("from-vault")
);
assert_eq!(PythonSecrets.get("ABSENT"), None);
assert_eq!(PythonSecrets.get("BROKEN"), None);
Python::attach(|py| {
let warnings: Vec<String> = py
.import("litellm.rust_bridge.settings")
.unwrap()
.getattr("warnings")
.unwrap()
.extract()
.unwrap();
assert_eq!(warnings.len(), 1);
assert!(warnings[0].contains("BROKEN") && warnings[0].contains("vault down"));
});
}
}

View file

@ -10,17 +10,16 @@ use litellm_auth_gcp::VertexAuth;
use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call};
use litellm_core::ocr::route::ocr_machine;
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings};
use litellm_llms::base_llm::ocr::{
handler::OcrClient,
settings::{OcrSettings, Secrets},
};
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use crate::{
errors::RustBridgeDeclined,
http,
python_settings::{PythonSecrets, PythonSettings},
};
use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSettings};
const SURFACE: LegacySurface = LegacySurface {
call_type: "ocr",
@ -42,6 +41,7 @@ fn run_ocr(
kwargs: Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let secrets = process_environment_secrets(&PythonSettings::SecretManager.read(py)?)?;
let config = http::call_config(py, &kwargs, asynchronous)?;
let client = OcrClient::new(
http::pool(),
@ -49,7 +49,7 @@ fn run_ocr(
http::url_policy(py)?,
VERTEX_AUTH.clone(),
ocr_settings(py)?,
Arc::new(PythonSecrets),
secrets,
)
.map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?;
run_legacy_call(
@ -62,6 +62,21 @@ fn run_ocr(
)
}
#[derive(FromPyObject)]
struct PythonSecretManager {
readable: bool,
}
fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult<Secrets> {
let manager: PythonSecretManager = secret_manager.extract()?;
if manager.readable {
return Err(RustBridgeDeclined::new_err(
"a readable secret manager is configured and the Rust route only reads the process environment",
));
}
Ok(Arc::new(ProcessEnvironment))
}
#[derive(FromPyObject)]
struct PythonProviderDefaults {
vertex_project: Option<String>,
@ -105,3 +120,47 @@ pub(crate) fn aocr(
) -> PyResult<Py<PyAny>> {
run_ocr(py, request, args, kwargs, true)
}
#[cfg(test)]
mod tests {
use pyo3::{prelude::*, types::PyDict};
use super::process_environment_secrets;
use crate::errors::RustBridgeDeclined;
fn secret_manager<'py>(py: Python<'py>, readable: bool) -> Bound<'py, PyAny> {
let locals = PyDict::new(py);
locals.set_item("readable", readable).unwrap();
py.run(
c"import types\nmanager = types.SimpleNamespace(readable=readable)",
Some(&locals),
Some(&locals),
)
.unwrap();
locals.get_item("manager").unwrap().unwrap()
}
#[test]
fn a_readable_secret_manager_sends_the_call_back_to_python() {
Python::initialize();
Python::attach(|py| {
let declined = process_environment_secrets(&secret_manager(py, true))
.err()
.expect("the Rust route declines");
assert!(declined.is_instance_of::<RustBridgeDeclined>(py));
});
}
#[test]
fn without_a_readable_secret_manager_secrets_are_the_process_environment() {
Python::initialize();
Python::attach(|py| {
let secrets = process_environment_secrets(&secret_manager(py, false)).unwrap();
assert_eq!(
secrets.get("LITELLM_RUST_BRIDGE_UNSET_VARIABLE_FOR_TEST"),
None
);
assert_eq!(secrets.get("PATH"), std::env::var("PATH").ok());
});
}
}

View file

@ -31,16 +31,21 @@ class ProviderDefaults:
enable_azure_ad_token_refresh: bool | None
@dataclass(frozen=True, slots=True)
class SecretManager:
readable: bool
def warn(message: str) -> None:
from litellm._logging import verbose_logger
verbose_logger.warning("%s", message)
def secret(name: str) -> str | None:
from litellm.secret_managers.main import get_secret_str
def secret_manager() -> SecretManager:
from litellm.secret_managers.main import _should_read_secret_from_secret_manager
return get_secret_str(name)
return SecretManager(readable=_should_read_secret_from_secret_manager())
def provider_defaults() -> ProviderDefaults:

View file

@ -11,6 +11,7 @@ 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 import settings
from litellm.secret_managers.main import get_secret_str
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json"
@ -23,6 +24,7 @@ def test_the_rust_contract_matches_the_returned_fields() -> None:
"http_settings": [field.name for field in dataclasses.fields(settings.http_settings())],
"url_policy": [field.name for field in dataclasses.fields(settings.url_policy())],
"provider_defaults": [field.name for field in dataclasses.fields(settings.provider_defaults())],
"secret_manager": [field.name for field in dataclasses.fields(settings.secret_manager())],
}
@ -101,25 +103,26 @@ class _VaultSecrets(CustomSecretManager):
return self.secrets.get(secret_name)
def test_secret_prefers_the_secret_manager_and_falls_back_to_the_environment_on_a_miss(
monkeypatch: pytest.MonkeyPatch,
@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", "stale-env-key")
monkeypatch.setenv("REDUCTO_API_KEY", "env-only-key")
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="read_only"))
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode))
assert settings.secret("MISTRAL_API_KEY") == "vault-key"
assert settings.secret("REDUCTO_API_KEY") == "env-only-key"
assert settings.secret("ABSENT_KEY") is None
assert settings.secret_manager() == settings.SecretManager(readable=readable)
assert (get_secret_str("MISTRAL_API_KEY") == "vault-key") is readable
def test_secret_reads_the_environment_without_a_secret_manager(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("MISTRAL_API_KEY", "env-key")
def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "secret_manager_client", None)
assert settings.secret("MISTRAL_API_KEY") == "env-key"
assert settings.secret_manager() == settings.SecretManager(readable=False)
def test_provider_defaults_read_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: