From 6c13e0912b43e54ba912f8257ce436d28e6e8e01 Mon Sep 17 00:00:00 2001 From: ryan Date: Sun, 20 Sep 2026 08:59:30 +0000 Subject: [PATCH 01/14] fix(proxy): apply DB-stored callback redaction settings before logger init Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 6 ++++ tests/test_litellm/proxy/test_proxy_server.py | 36 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/litellm/constants.py b/litellm/constants.py index bbeb4846e27..4bc2e4e0d26 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1853,6 +1853,12 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ "max_ui_session_budget", "budget_rollover", "mcp_tool_search", + "turn_off_message_logging", + "datadog_params", + "datadog_llm_observability_params", + "newrelic_params", + "pointfive_params", + "aws_sqs_callback_params", ] SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index f71f9c20f3b..75741b125cb 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11766,6 +11766,42 @@ def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_n assert getattr(litellm, field_name) == db_value +@pytest.mark.asyncio +async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(monkeypatch): + """A DB-only litellm_settings row that pairs success_callback: ["datadog"] with + datadog_params.turn_off_message_logging: true must build the DataDogLogger redacted, the + same as the identical block in YAML. Regression for the redaction keys being absent from + the safe-override allowlist while the callback half of the row was honoured.""" + import litellm.proxy.proxy_server as ps + from litellm.integrations.datadog.datadog import DataDogLogger + from litellm.litellm_core_utils import litellm_logging + + monkeypatch.setenv("DD_API_KEY", "test-key") + monkeypatch.setenv("DD_SITE", "us5.datadoghq.com") + monkeypatch.setattr(litellm, "datadog_params", None) + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", []) + + db_row = { + "success_callback": ["datadog"], + "datadog_params": {"turn_off_message_logging": True}, + "turn_off_message_logging": True, + } + pc = ps.ProxyConfig() + pc._apply_litellm_settings_db_values(pc._prepared_db_settings_values("litellm_settings", db_row)) + pc._add_callbacks_from_db_config({"litellm_settings": db_row}) + + datadog_loggers = [cb for cb in litellm.success_callback if isinstance(cb, DataDogLogger)] + assert len(datadog_loggers) == 1 + assert datadog_loggers[0].turn_off_message_logging is True + assert litellm.turn_off_message_logging is True + + def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypatch): """The flag defaults to False rather than None, so a plain 'is not None' check would report the default as 'In Config' and imply an admin had set it.""" From a90d852d378b46e3541fbb76ec519845e5caa785 Mon Sep 17 00:00:00 2001 From: ryan Date: Sun, 20 Sep 2026 09:26:30 +0000 Subject: [PATCH 02/14] test(proxy): type monkeypatch and cover every DB-overridable callback params key Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_server.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 75741b125cb..c56645313c3 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11767,7 +11767,7 @@ def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_n @pytest.mark.asyncio -async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(monkeypatch): +async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(monkeypatch: pytest.MonkeyPatch): """A DB-only litellm_settings row that pairs success_callback: ["datadog"] with datadog_params.turn_off_message_logging: true must build the DataDogLogger redacted, the same as the identical block in YAML. Regression for the redaction keys being absent from @@ -11802,6 +11802,30 @@ async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(mon assert litellm.turn_off_message_logging is True +@pytest.mark.parametrize( + "field_name", + [ + "datadog_params", + "datadog_llm_observability_params", + "newrelic_params", + "pointfive_params", + "aws_sqs_callback_params", + ], +) +def test_db_stored_callback_params_propagate_to_litellm_module(monkeypatch: pytest.MonkeyPatch, field_name: str): + """Every callback init params block stored in the DB litellm_settings row must land on the + litellm module before the matching logger is built, so the DB row behaves like YAML.""" + import litellm.proxy.proxy_server as ps + + monkeypatch.setattr(litellm, field_name, None) + db_value = {"turn_off_message_logging": True} + + pc = ps.ProxyConfig() + pc._apply_litellm_settings_db_values(pc._prepared_db_settings_values("litellm_settings", {field_name: db_value})) + + assert getattr(litellm, field_name) == db_value + + def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypatch): """The flag defaults to False rather than None, so a plain 'is not None' check would report the default as 'In Config' and imply an admin had set it.""" From 3ba4a60d5ed5eeeb217b63ac7898746927f0db12 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:26:28 +0000 Subject: [PATCH 03/14] feat(rust): add HashiCorp Vault secret manager crate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-rust.yml | 2 +- litellm-rust/Cargo.lock | 19 + litellm-rust/Cargo.toml | 1 + .../crates/secrets-hashicorp/Cargo.toml | 23 + .../crates/secrets-hashicorp/src/config.rs | 169 +++++++ .../crates/secrets-hashicorp/src/error.rs | 32 ++ .../crates/secrets-hashicorp/src/lib.rs | 9 + .../secrets-hashicorp/src/secret_manager.rs | 333 ++++++++++++++ .../secrets-hashicorp/tests/secret_manager.rs | 424 ++++++++++++++++++ litellm-rust/crates/secrets/Cargo.toml | 2 + litellm-rust/crates/secrets/README.md | 2 + litellm-rust/crates/secrets/src/error.rs | 3 + litellm-rust/crates/secrets/src/handler.rs | 10 + litellm-rust/crates/secrets/src/lib.rs | 2 + litellm-rust/crates/secrets/tests/handler.rs | 128 ++++++ .../hashicorp_secret_manager.py | 16 +- .../hashicorp_vault_parity.json | 97 ++++ .../test_hashicorp_secret_manager.py | 50 +++ 18 files changed, 1314 insertions(+), 8 deletions(-) create mode 100644 litellm-rust/crates/secrets-hashicorp/Cargo.toml create mode 100644 litellm-rust/crates/secrets-hashicorp/src/config.rs create mode 100644 litellm-rust/crates/secrets-hashicorp/src/error.rs create mode 100644 litellm-rust/crates/secrets-hashicorp/src/lib.rs create mode 100644 litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs create mode 100644 tests/test_litellm/secret_managers/hashicorp_vault_parity.json diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 278fa7c425f..56bb9a568fc 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -130,7 +130,7 @@ jobs: - name: Test secret manager feature combinations run: | cargo test -p litellm-auth-gcp --locked --no-default-features - for features in '' aws google aws,google; do + for features in '' aws google hashicorp aws,google,hashicorp; do cargo test -p litellm-secrets --locked --no-default-features --features "$features" done diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..6f05bba9417 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2698,6 +2698,7 @@ dependencies = [ "litellm-core-utils", "litellm-secrets-aws", "litellm-secrets-google", + "litellm-secrets-hashicorp", "litellm-secrets-types", "moka", "reqwest 0.12.28", @@ -2755,6 +2756,24 @@ dependencies = [ "wiremock", ] +[[package]] +name = "litellm-secrets-hashicorp" +version = "0.1.0" +dependencies = [ + "litellm-core-utils", + "litellm-secrets-types", + "moka", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.19", + "tokio", + "veil", + "wiremock", +] + [[package]] name = "litellm-secrets-types" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..5030a94b140 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -22,6 +22,7 @@ litellm-secrets = { path = "crates/secrets" } litellm-secrets-types = { path = "crates/secrets-types" } litellm-secrets-aws = { path = "crates/secrets-aws" } litellm-secrets-google = { path = "crates/secrets-google" } +litellm-secrets-hashicorp = { path = "crates/secrets-hashicorp" } litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml new file mode 100644 index 00000000000..0656980a033 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "litellm-secrets-hashicorp" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-core-utils.workspace = true +litellm-secrets-types.workspace = true +moka.workspace = true +reqwest.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true +tempfile = "3" +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs new file mode 100644 index 00000000000..f9491f71afb --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -0,0 +1,169 @@ +use std::{path::PathBuf, time::Duration}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::KeyManagementSettings; +use litellm_secrets_types::SecretValue; + +use crate::Error; + +const DEFAULT_ADDRESS: &str = "http://127.0.0.1:8200"; +const DEFAULT_MOUNT: &str = "secret"; +const DEFAULT_APPROLE_MOUNT_PATH: &str = "approle"; +const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(86400); +const HCP_VAULT_ADDR: &str = "HCP_VAULT_ADDR"; +const HCP_VAULT_TOKEN: &str = "HCP_VAULT_TOKEN"; +const HCP_VAULT_NAMESPACE: &str = "HCP_VAULT_NAMESPACE"; +const HCP_VAULT_LOGIN_NAMESPACE: &str = "HCP_VAULT_LOGIN_NAMESPACE"; +const HCP_VAULT_SECRET_NAMESPACE: &str = "HCP_VAULT_SECRET_NAMESPACE"; +const HCP_VAULT_MOUNT_NAME: &str = "HCP_VAULT_MOUNT_NAME"; +const HCP_VAULT_PATH_PREFIX: &str = "HCP_VAULT_PATH_PREFIX"; +const HCP_VAULT_APPROLE_ROLE_ID: &str = "HCP_VAULT_APPROLE_ROLE_ID"; +const HCP_VAULT_APPROLE_SECRET_ID: &str = "HCP_VAULT_APPROLE_SECRET_ID"; +const HCP_VAULT_APPROLE_MOUNT_PATH: &str = "HCP_VAULT_APPROLE_MOUNT_PATH"; +const HCP_VAULT_CLIENT_CERT: &str = "HCP_VAULT_CLIENT_CERT"; +const HCP_VAULT_CLIENT_KEY: &str = "HCP_VAULT_CLIENT_KEY"; +const HCP_VAULT_CERT_ROLE: &str = "HCP_VAULT_CERT_ROLE"; +const HCP_VAULT_REFRESH_INTERVAL: &str = "HCP_VAULT_REFRESH_INTERVAL"; +const SECRET_MANAGER_REFRESH_INTERVAL: &str = "SECRET_MANAGER_REFRESH_INTERVAL"; + +#[derive(Clone, Debug)] +pub struct AppRoleAuth { + pub role_id: String, + pub secret_id: SecretValue, + pub mount_path: String, +} + +#[derive(Clone, Debug)] +pub struct TlsCertAuth { + pub cert_path: PathBuf, + pub key_path: PathBuf, + pub role: Option, +} + +#[derive(Clone, Debug)] +pub struct HashicorpVaultConfig { + pub address: String, + pub token: Option, + pub namespace: Option, + pub login_namespace: Option, + pub secret_namespace: Option, + pub mount: String, + pub path_prefix: Option, + pub approle: Option, + pub tls_cert: Option, + pub refresh_interval: Duration, +} + +impl HashicorpVaultConfig { + pub fn from_environment(environment: &dyn Lookup) -> Result { + let address: String = environment + .get(HCP_VAULT_ADDR) + .and_then(|value| nonempty(value.trim())) + .map(|value| value.trim_end_matches('/').to_owned()) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| DEFAULT_ADDRESS.to_owned()); + let token: Option = environment + .get(HCP_VAULT_TOKEN) + .and_then(nonempty) + .map(SecretValue::new); + let namespace: Option = path_component(environment.get(HCP_VAULT_NAMESPACE)); + let login_namespace: Option = + path_component(environment.get(HCP_VAULT_LOGIN_NAMESPACE)); + let secret_namespace: Option = + path_component(environment.get(HCP_VAULT_SECRET_NAMESPACE)); + let mount: String = path_component(environment.get(HCP_VAULT_MOUNT_NAME)) + .unwrap_or_else(|| DEFAULT_MOUNT.to_owned()); + let path_prefix: Option = path_component(environment.get(HCP_VAULT_PATH_PREFIX)); + let approle: Option = match ( + environment + .get(HCP_VAULT_APPROLE_ROLE_ID) + .and_then(nonempty), + environment + .get(HCP_VAULT_APPROLE_SECRET_ID) + .and_then(nonempty) + .map(SecretValue::new), + ) { + (Some(role_id), Some(secret_id)) => Some(AppRoleAuth { + role_id, + secret_id, + mount_path: path_component(environment.get(HCP_VAULT_APPROLE_MOUNT_PATH)) + .unwrap_or_else(|| DEFAULT_APPROLE_MOUNT_PATH.to_owned()), + }), + _ => None, + }; + let tls_cert: Option = match ( + environment.get(HCP_VAULT_CLIENT_CERT).and_then(nonempty), + environment.get(HCP_VAULT_CLIENT_KEY).and_then(nonempty), + ) { + (Some(cert_path), Some(key_path)) => Some(TlsCertAuth { + cert_path: PathBuf::from(cert_path), + key_path: PathBuf::from(key_path), + role: environment.get(HCP_VAULT_CERT_ROLE).and_then(nonempty), + }), + _ => None, + }; + let refresh_interval: Duration = refresh_interval(environment)?; + Ok(Self { + address, + token, + namespace, + login_namespace, + secret_namespace, + mount, + path_prefix, + approle, + tls_cert, + refresh_interval, + }) + } + + pub fn from_settings( + _settings: &KeyManagementSettings, + environment: &dyn Lookup, + ) -> Result { + Self::from_environment(environment) + } + + pub fn login_namespace(&self) -> Option<&str> { + self.login_namespace + .as_deref() + .or(self.namespace.as_deref()) + } + + pub fn secret_namespace(&self) -> Option<&str> { + self.secret_namespace + .as_deref() + .or(self.namespace.as_deref()) + } +} + +fn nonempty(value: impl AsRef) -> Option { + let value: &str = value.as_ref(); + (!value.is_empty()).then(|| value.to_owned()) +} + +fn path_component(value: Option) -> Option { + value + .and_then(|value| nonempty(value.trim())) + .map(|value| value.trim_matches('/').to_owned()) + .filter(|value| !value.is_empty()) +} + +fn refresh_interval(environment: &dyn Lookup) -> Result { + let value: Option = environment + .get(HCP_VAULT_REFRESH_INTERVAL) + .and_then(nonempty) + .or_else(|| { + environment + .get(SECRET_MANAGER_REFRESH_INTERVAL) + .and_then(nonempty) + }); + let Some(value) = value else { + return Ok(DEFAULT_REFRESH_INTERVAL); + }; + let seconds: i64 = value.parse().map_err(|_| Error::RefreshInterval)?; + if seconds < 0 { + return Ok(Duration::from_nanos(1)); + } + Ok(Duration::from_secs(seconds as u64)) +} diff --git a/litellm-rust/crates/secrets-hashicorp/src/error.rs b/litellm-rust/crates/secrets-hashicorp/src/error.rs new file mode 100644 index 00000000000..3f6a5b66475 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/error.rs @@ -0,0 +1,32 @@ +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("HashiCorp Vault requires an enterprise license")] + EnterpriseRequired, + #[error("invalid secret name")] + InvalidSecretName(#[from] litellm_secrets_types::Error), + #[error("HashiCorp Vault request failed")] + Request( + #[from] + #[redact] + reqwest::Error, + ), + #[error("HashiCorp Vault TLS identity could not be configured for {path}: {message}")] + TlsIdentity { + path: std::path::PathBuf, + message: String, + }, + #[error("HashiCorp Vault login returned HTTP {status}")] + LoginStatus { status: u16 }, + #[error("HashiCorp Vault login response is malformed")] + MalformedLogin, + #[error("HashiCorp Vault authentication is not configured")] + NoAuthConfigured, + #[error("HashiCorp Vault returned HTTP {status}")] + Status { status: u16 }, + #[error("HashiCorp Vault response payload is malformed")] + MalformedPayload, + #[error("HashiCorp Vault secret value is not a string")] + NonStringValue, + #[error("invalid HashiCorp Vault refresh interval")] + RefreshInterval, +} diff --git a/litellm-rust/crates/secrets-hashicorp/src/lib.rs b/litellm-rust/crates/secrets-hashicorp/src/lib.rs new file mode 100644 index 00000000000..a561a58b59e --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/lib.rs @@ -0,0 +1,9 @@ +#![forbid(unsafe_code)] + +mod config; +mod error; +pub mod secret_manager; + +pub use config::{AppRoleAuth, HashicorpVaultConfig, TlsCertAuth}; +pub use error::Error; +pub use secret_manager::HashicorpVault; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs new file mode 100644 index 00000000000..6b887289443 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -0,0 +1,333 @@ +use std::{ + fmt, + sync::Arc, + time::{Duration, Instant}, +}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{ + BaseSecretManager, SecretValue, async_rotate_secret, validate_secret_name, +}; +use moka::future::Cache; +use reqwest::Client; +use serde_json::{Value, json}; +use tokio::sync::Mutex; + +use crate::{Error, HashicorpVaultConfig}; + +const CACHE_CAPACITY: u64 = 200; + +#[derive(Clone)] +struct CachedToken { + token: SecretValue, + expires_at: Option, +} + +#[derive(Clone)] +pub struct HashicorpVault { + client: Arc, + config: HashicorpVaultConfig, + cache: Cache, + auth_token: Arc>>, +} + +impl fmt::Debug for HashicorpVault { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("HashicorpVault") + .field("config", &self.config) + .finish_non_exhaustive() + } +} + +impl HashicorpVault { + pub fn new( + environment: Arc, + enterprise_enabled: bool, + ) -> Result { + if !enterprise_enabled { + return Err(Error::EnterpriseRequired); + } + let config: HashicorpVaultConfig = + HashicorpVaultConfig::from_environment(environment.as_ref())?; + let client: Client = client_for_config(&config)?; + Self::with_client(client, config, enterprise_enabled) + } + + pub fn with_client( + client: Client, + config: HashicorpVaultConfig, + enterprise_enabled: bool, + ) -> Result { + if !enterprise_enabled { + return Err(Error::EnterpriseRequired); + } + let cache: Cache = Cache::builder() + .max_capacity(CACHE_CAPACITY) + .time_to_live(config.refresh_interval) + .build(); + Ok(Self { + client: Arc::new(client), + config, + cache, + auth_token: Arc::new(Mutex::new(None)), + }) + } + + pub fn secret_url(&self, secret_name: &str) -> Result { + validate_secret_name(secret_name).map_err(Error::InvalidSecretName)?; + let namespace: String = self + .config + .secret_namespace() + .map(|value| format!("{value}/")) + .unwrap_or_default(); + let path_prefix: String = self + .config + .path_prefix + .as_deref() + .map(|value| format!("{value}/")) + .unwrap_or_default(); + Ok(format!( + "{}/v1/{}{}/data/{}{}", + self.config.address, namespace, self.config.mount, path_prefix, secret_name + )) + } + + pub fn login_url(&self) -> Option { + self.config.approle.as_ref().map_or_else( + || { + self.config + .tls_cert + .as_ref() + .map(|_| format!("{}/v1/auth/cert/login", self.config.address)) + }, + |approle| { + Some(format!( + "{}/v1/auth/{}/login", + self.config.address, approle.mount_path + )) + }, + ) + } + + pub fn config(&self) -> &HashicorpVaultConfig { + &self.config + } + + pub async fn async_read_secret(&self, secret_name: &str) -> Result, Error> { + let url: String = self.secret_url(secret_name)?; + if let Some(value) = self.cache.get(&url).await { + return Ok(Some(value)); + } + let token: SecretValue = self.vault_token().await?; + let mut request = self + .client + .get(&url) + .header("X-Vault-Token", token.expose()); + if let Some(namespace) = self.config.secret_namespace() { + request = request.header("X-Vault-Namespace", namespace); + } + let response: reqwest::Response = request.send().await?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + if !response.status().is_success() { + return Err(Error::Status { + status: response.status().as_u16(), + }); + } + let body = response.bytes().await?; + let body: Value = serde_json::from_slice(&body).map_err(|_| Error::MalformedPayload)?; + let data: &Value = body + .get("data") + .and_then(Value::as_object) + .and_then(|value| value.get("data")) + .ok_or(Error::MalformedPayload)?; + let data: &serde_json::Map = + data.as_object().ok_or(Error::MalformedPayload)?; + let Some(value) = data.get("key") else { + return Ok(None); + }; + let value: &str = value.as_str().ok_or(Error::NonStringValue)?; + let value: SecretValue = SecretValue::new(value); + self.cache.insert(url, value.clone()).await; + Ok(Some(value)) + } + + pub async fn async_write_secret( + &self, + secret_name: &str, + value: SecretValue, + description: Option<&str>, + ) -> Result { + let url: String = self.secret_url(secret_name)?; + let data: Value = match description { + Some(description) => json!({"key": value.expose(), "description": description}), + None => json!({"key": value.expose()}), + }; + let token: SecretValue = self.vault_token().await?; + let mut request = self + .client + .post(&url) + .header("X-Vault-Token", token.expose()); + if let Some(namespace) = self.config.secret_namespace() { + request = request.header("X-Vault-Namespace", namespace); + } + let response: reqwest::Response = request.json(&json!({"data": data})).send().await?; + if !response.status().is_success() { + return Err(Error::Status { + status: response.status().as_u16(), + }); + } + self.cache.invalidate(&url).await; + let body = response.bytes().await?; + if body.is_empty() { + return Ok(Value::Null); + } + serde_json::from_slice(&body).map_err(|_| Error::MalformedPayload) + } + + pub async fn async_delete_secret(&self, secret_name: &str) -> Result<(), Error> { + let url: String = self.secret_url(secret_name)?; + let token: SecretValue = self.vault_token().await?; + let mut request = self + .client + .delete(&url) + .header("X-Vault-Token", token.expose()); + if let Some(namespace) = self.config.secret_namespace() { + request = request.header("X-Vault-Namespace", namespace); + } + let response: reqwest::Response = request.send().await?; + if !response.status().is_success() { + return Err(Error::Status { + status: response.status().as_u16(), + }); + } + self.cache.invalidate(&url).await; + Ok(()) + } + + pub async fn async_rotate_secret( + &self, + current_name: &str, + new_name: &str, + value: &SecretValue, + ) -> Result { + async_rotate_secret(self, current_name, new_name, value).await + } + + async fn vault_token(&self) -> Result { + let mut cached: tokio::sync::MutexGuard<'_, Option> = + self.auth_token.lock().await; + if let Some(entry) = cached.as_ref() + && entry + .expires_at + .is_none_or(|expires_at| expires_at > Instant::now()) + { + return Ok(entry.token.clone()); + } + let Some(login_url) = self.login_url() else { + let Some(token) = self.config.token.clone() else { + return Err(Error::NoAuthConfigured); + }; + return Ok(token); + }; + let body: Value = match (self.config.approle.as_ref(), self.config.tls_cert.as_ref()) { + (Some(approle), _) => { + json!({"role_id": approle.role_id, "secret_id": approle.secret_id.expose()}) + } + (None, Some(tls)) => tls + .role + .as_deref() + .map_or_else(|| json!({}), |role| json!({"name": role})), + (None, None) => { + let Some(token) = self.config.token.clone() else { + return Err(Error::NoAuthConfigured); + }; + return Ok(token); + } + }; + let mut request = self.client.post(login_url).json(&body); + if let Some(namespace) = self.config.login_namespace() { + request = request.header("X-Vault-Namespace", namespace); + } + let response: reqwest::Response = request.send().await?; + if !response.status().is_success() { + return Err(Error::LoginStatus { + status: response.status().as_u16(), + }); + } + let body = response.bytes().await?; + let payload: Value = serde_json::from_slice(&body).map_err(|_| Error::MalformedLogin)?; + let auth: &serde_json::Map = payload + .get("auth") + .and_then(Value::as_object) + .ok_or(Error::MalformedLogin)?; + let token: SecretValue = SecretValue::new( + auth.get("client_token") + .and_then(Value::as_str) + .ok_or(Error::MalformedLogin)?, + ); + let lease_duration: u64 = auth + .get("lease_duration") + .and_then(Value::as_u64) + .ok_or(Error::MalformedLogin)?; + let expires_at: Option = + (lease_duration > 0).then(|| Instant::now() + Duration::from_secs(lease_duration)); + *cached = Some(CachedToken { + token: token.clone(), + expires_at, + }); + Ok(token) + } +} + +impl BaseSecretManager for HashicorpVault { + type Error = Error; + type WriteResponse = Value; + type DeleteResponse = (); + + async fn async_read_secret(&self, name: &str) -> Result, Error> { + HashicorpVault::async_read_secret(self, name).await + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + HashicorpVault::async_write_secret(self, name, value.clone(), description).await + } + + async fn async_delete_secret( + &self, + name: &str, + _recovery_window_in_days: i64, + ) -> Result<(), Error> { + HashicorpVault::async_delete_secret(self, name).await + } +} + +fn client_for_config(config: &HashicorpVaultConfig) -> Result { + let mut builder: reqwest::ClientBuilder = Client::builder(); + if let Some(tls) = config.tls_cert.as_ref() { + let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + })?; + let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { + path: tls.key_path.clone(), + message: source.to_string(), + })?; + let identity: reqwest::Identity = reqwest::Identity::from_pem( + &[cert.as_slice(), key.as_slice()].concat(), + ) + .map_err(|source| Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + })?; + builder = builder.identity(identity); + } + builder.build().map_err(Error::Request) +} diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs new file mode 100644 index 00000000000..1cf13a79c77 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -0,0 +1,424 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets_hashicorp::{Error, HashicorpVault, HashicorpVaultConfig}; +use litellm_secrets_types::SecretValue; +use serde::Deserialize; +use serde_json::json; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, header, method, path}, +}; + +fn config(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVaultConfig { + let mut environment_values: HashMap = values + .iter() + .map(|(name, value)| ((*name).to_owned(), (*value).to_owned())) + .collect(); + environment_values.insert("HCP_VAULT_ADDR".to_owned(), server.uri()); + let environment: Arc = + Arc::new(move |name: &str| environment_values.get(name).cloned()); + HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap() +} + +fn manager(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVault { + HashicorpVault::with_client(reqwest::Client::new(), config(server, values), true).unwrap() +} + +#[tokio::test] +async fn token_reads_use_vault_headers_and_cache_values() { + let server: MockServer = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/name")) + .and(header("X-Vault-Token", "token")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(1) + .mount(&server) + .await; + let manager: HashicorpVault = manager(&server, &[("HCP_VAULT_TOKEN", "token")]); + + assert_eq!( + manager + .async_read_secret("name") + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); + assert_eq!( + manager + .async_read_secret("name") + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} + +#[tokio::test] +async fn namespace_mount_and_prefix_are_sanitized_in_the_url() { + let server: MockServer = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/team-a/kv-prod/data/virtual-keys/name")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(1) + .mount(&server) + .await; + let manager: HashicorpVault = manager( + &server, + &[ + ("HCP_VAULT_TOKEN", "token"), + ("HCP_VAULT_SECRET_NAMESPACE", " /team-a/ "), + ("HCP_VAULT_MOUNT_NAME", " /kv-prod/ "), + ("HCP_VAULT_PATH_PREFIX", " /virtual-keys/ "), + ], + ); + + assert_eq!( + manager.secret_url("name").unwrap(), + format!("{}/v1/team-a/kv-prod/data/virtual-keys/name", server.uri()) + ); + assert!(manager.async_read_secret("name").await.unwrap().is_some()); +} + +#[tokio::test] +async fn approle_login_uses_namespace_and_reuses_the_token() { + let server: MockServer = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/auth/custom-approle/login")) + .and(header("X-Vault-Namespace", "login-root")) + .and(body_json(json!({"role_id": "role", "secret_id": "secret"}))) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"auth": {"client_token": "login-token", "lease_duration": 3600}}), + )) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/secret-root/secret/data/name")) + .and(header("X-Vault-Token", "login-token")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(1) + .mount(&server) + .await; + let manager: HashicorpVault = manager( + &server, + &[ + ("HCP_VAULT_APPROLE_ROLE_ID", "role"), + ("HCP_VAULT_APPROLE_SECRET_ID", "secret"), + ("HCP_VAULT_APPROLE_MOUNT_PATH", "custom-approle"), + ("HCP_VAULT_NAMESPACE", "secret-root"), + ("HCP_VAULT_LOGIN_NAMESPACE", "login-root"), + ], + ); + + assert!(manager.async_read_secret("name").await.unwrap().is_some()); + assert!(manager.async_read_secret("name-2").await.unwrap().is_none()); +} + +#[tokio::test] +async fn approle_tokens_expire_after_the_vault_lease() { + let server: MockServer = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/auth/approle/login")) + .respond_with( + ResponseTemplate::new(200).set_body_json( + json!({"auth": {"client_token": "login-token", "lease_duration": 1}}), + ), + ) + .expect(2) + .mount(&server) + .await; + Mock::given(method("GET")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(2) + .mount(&server) + .await; + let manager: HashicorpVault = manager( + &server, + &[ + ("HCP_VAULT_APPROLE_ROLE_ID", "role"), + ("HCP_VAULT_APPROLE_SECRET_ID", "secret"), + ("HCP_VAULT_REFRESH_INTERVAL", "0"), + ], + ); + + assert!(manager.async_read_secret("first").await.unwrap().is_some()); + tokio::time::sleep(Duration::from_secs(1) + Duration::from_millis(50)).await; + assert!(manager.async_read_secret("second").await.unwrap().is_some()); +} + +#[tokio::test] +async fn tls_login_posts_the_role_and_uses_the_client_identity() { + let server: MockServer = MockServer::start().await; + let directory: tempfile::TempDir = tempfile::tempdir().unwrap(); + let cert_path = directory.path().join("client.crt"); + let key_path = directory.path().join("client.key"); + std::fs::write(&cert_path, TEST_CERTIFICATE).unwrap(); + std::fs::write(&key_path, TEST_PRIVATE_KEY).unwrap(); + Mock::given(method("POST")) + .and(path("/v1/auth/cert/login")) + .and(body_json(json!({"name": "vault-role"}))) + .respond_with( + ResponseTemplate::new(200).set_body_json( + json!({"auth": {"client_token": "cert-token", "lease_duration": 0}}), + ), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/name")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .mount(&server) + .await; + let manager: HashicorpVault = HashicorpVault::new( + { + let environment_values: HashMap = HashMap::from([ + ("HCP_VAULT_ADDR".to_owned(), server.uri()), + ( + "HCP_VAULT_CLIENT_CERT".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_CLIENT_KEY".to_owned(), + key_path.to_str().unwrap().to_owned(), + ), + ("HCP_VAULT_CERT_ROLE".to_owned(), "vault-role".to_owned()), + ]); + Arc::new(move |name: &str| environment_values.get(name).cloned()) + }, + true, + ) + .unwrap(); + + assert!(manager.async_read_secret("name").await.unwrap().is_some()); + assert_eq!( + manager.login_url().as_deref(), + Some(format!("{}/v1/auth/cert/login", server.uri()).as_str()) + ); +} + +#[rstest::rstest] +#[case::missing(404, json!({}), 0)] +#[case::malformed(200, json!({}), 1)] +#[case::missing_key(200, json!({"data": {"data": {}}}), 0)] +#[case::non_string(200, json!({"data": {"data": {"key": 1}}}), 2)] +#[tokio::test] +async fn read_responses_distinguish_absence_and_malformed_payloads( + #[case] status: u16, + #[case] body: serde_json::Value, + #[case] expected: u8, +) { + let server: MockServer = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(status).set_body_json(body)) + .expect(1) + .mount(&server) + .await; + let result: Result, Error> = + manager(&server, &[("HCP_VAULT_TOKEN", "token")]) + .async_read_secret("name") + .await; + match expected { + 0 => assert!(result.unwrap().is_none()), + 1 => assert!(matches!(result, Err(Error::MalformedPayload))), + 2 => assert!(matches!(result, Err(Error::NonStringValue))), + _ => unreachable!(), + } +} + +#[tokio::test] +async fn write_and_delete_invalidate_the_read_cache() { + let server: MockServer = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/name")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ) + .expect(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/secret/data/name")) + .and(body_json( + json!({"data": {"key": "updated", "description": "description"}}), + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data": {"version": 2}}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("DELETE")) + .and(path("/v1/secret/data/name")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let manager: HashicorpVault = manager(&server, &[("HCP_VAULT_TOKEN", "token")]); + + assert!(manager.async_read_secret("name").await.unwrap().is_some()); + assert!( + manager + .async_write_secret("name", SecretValue::new("updated"), Some("description")) + .await + .is_ok() + ); + assert!(manager.async_read_secret("name").await.unwrap().is_some()); + manager.async_delete_secret("name").await.unwrap(); +} + +#[tokio::test] +async fn no_auth_and_invalid_names_fail_without_requests() { + let server: MockServer = MockServer::start().await; + let manager: HashicorpVault = manager(&server, &[]); + + assert!(matches!( + manager.async_read_secret("name").await, + Err(Error::NoAuthConfigured) + )); + assert!(matches!( + manager.async_read_secret("../name").await, + Err(Error::InvalidSecretName(_)) + )); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn debug_output_redacts_authentication_values() { + let server: MockServer = MockServer::start().await; + let manager: HashicorpVault = HashicorpVault::with_client( + reqwest::Client::new(), + config(&server, &[("HCP_VAULT_TOKEN", "token-value")]), + true, + ) + .unwrap(); + let debug: String = format!("{manager:?}"); + assert!(!debug.contains("token-value")); + assert!(!debug.contains("secret-id")); +} + +#[derive(Deserialize)] +struct ParityCase { + env: HashMap, + expected_secret_url: String, + expected_login_url: Option, + expected_login_namespace: Option, + expected_secret_namespace: Option, + secret_name: String, +} + +#[test] +fn configuration_matches_python_parity_fixture() { + let cases: Vec = serde_json::from_str(include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../tests/test_litellm/secret_managers/hashicorp_vault_parity.json" + ))) + .unwrap(); + for case in cases { + let values: HashMap = case.env.clone(); + let environment: Arc = + Arc::new(move |name: &str| values.get(name).cloned()); + let config: HashicorpVaultConfig = + HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap(); + let manager: HashicorpVault = + HashicorpVault::with_client(reqwest::Client::new(), config, true).unwrap(); + assert_eq!( + manager.secret_url(&case.secret_name).unwrap(), + case.expected_secret_url + ); + assert_eq!(manager.login_url(), case.expected_login_url); + assert_eq!( + manager.config().login_namespace(), + case.expected_login_namespace.as_deref() + ); + assert_eq!( + manager.config().secret_namespace(), + case.expected_secret_namespace.as_deref() + ); + } +} + +#[tokio::test] +#[ignore] +async fn live_vault_round_trip() { + let environment: Arc = + Arc::new(litellm_core_utils::settings::ProcessEnvironment); + let manager: HashicorpVault = HashicorpVault::new(environment, true).unwrap(); + let name: String = std::env::var("LITELLM_VAULT_LIVE_SECRET_NAME").unwrap(); + let value: SecretValue = SecretValue::new("native-live-value"); + let url: String = manager.secret_url(&name).unwrap(); + println!("native provenance: {} {}", module_path!(), url); + manager + .async_write_secret(&name, value.clone(), None) + .await + .unwrap(); + assert_eq!( + manager.async_read_secret(&name).await.unwrap().unwrap(), + value + ); + manager.async_delete_secret(&name).await.unwrap(); + assert!(manager.async_read_secret(&name).await.unwrap().is_none()); +} + +const TEST_CERTIFICATE: &str = "-----BEGIN CERTIFICATE----- +MIIDDzCCAfegAwIBAgIUeMzLFLM/mRbPGbNAew5N2UTscocwDQYJKoZIhvcNAQEL +BQAwFzEVMBMGA1UEAwwMbGl0ZWxsbS10ZXN0MB4XDTI2MDkyMTIwMjA1OVoXDTI2 +MDkyMjIwMjA1OVowFzEVMBMGA1UEAwwMbGl0ZWxsbS10ZXN0MIIBIjANBgkqhkiG +9w0BAQEFAAOCAQ8AMIIBCgKCAQEAveYoSUJXybmkHmQsBfhBcv2Ob5Oy8ejZu+B3 +vTnrPumW4ANi1XXKBSazRGB3fEtAgr+3KhKeHaSKEQeBwJkAEBfdmQv0tpXICwHs +1kFNtU0owy54HVW5/ia+LMszsFcPzVIoMnbUOuiKr9RaV7P+IEFzILPBVuV4DoYH +yocjD3+9QNqokWgNL8LK37JijmNEFVaKFz0X6SyL2VRDlfPWTEBK52Gp/pvDgA6G +eTSfyI+kCm9h5ECTYUAtmatk9WPVS8sWOqV1EXVanFyYBU+mDxoywAS1/6CHeIPh +bNmCOZjPoO9qWBJ7ZyGhOconBigXY8qnlXymev+44IPHrx4urwIDAQABo1MwUTAd +BgNVHQ4EFgQUvaZrZ6HKtbr3ekeZmgy4b5Pq95QwHwYDVR0jBBgwFoAUvaZrZ6HK +tbr3ekeZmgy4b5Pq95QwDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0BAQsFAAOC +AQEAEejrD8d1qDxW55XxQ4IC31rufoEvDV955jyvh2kALPaN/i5oWsBGI+UAQZna +aaoQXwzlmHrtDUBWl0LztVTUamIleUep2+PLLauqqt43vxppxMX8Jn2mnPO20YE/ +hIzGx0jN/LBG8PDyLSvHdlgjP9ofA4Vg4rTQugdXRgOvlCE/epnH/MADcg9KYJtJ +C1RObCIkL3LcdUbjStJRCY/U/FeWcgyncEPz95OFDkbrlNDajb6o6CkYfouqvhTc +8XlgjjAVKIbAbRgbVu3elsquuFM97x2DzWDjkrMNmDt1FJ9ubK36gL6B3o0UMaoQ +00R7x/eqvH+EkWa/2ekW9lpleQ== +-----END CERTIFICATE----- +"; + +const TEST_PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY----- +MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC95ihJQlfJuaQe +ZCwF+EFy/Y5vk7Lx6Nm74He9Oes+6ZbgA2LVdcoFJrNEYHd8S0CCv7cqEp4dpIoR +B4HAmQAQF92ZC/S2lcgLAezWQU21TSjDLngdVbn+Jr4syzOwVw/NUigydtQ66Iqv +1FpXs/4gQXMgs8FW5XgOhgfKhyMPf71A2qiRaA0vwsrfsmKOY0QVVooXPRfpLIvZ +VEOV89ZMQErnYan+m8OADoZ5NJ/Ij6QKb2HkQJNhQC2Zq2T1Y9VLyxY6pXURdVqc +XJgFT6YPGjLABLX/oId4g+Fs2YI5mM+g72pYEntnIaE5yicGKBdjyqeVfKZ6/7jg +g8evHi6vAgMBAAECggEAGdJjlP6b8Fa5bdaCM/ebcrbuuNZVJVbb0JPHxGfNSLs7 +pE9hj5QaOdQW2Uviw3h6F61ZCzQH4xD+Iy2po5ZKb2XHYKnDB1bboj+LRGER337T +9aJqe9at2VTMVEv3Rdm40NsEk0QcPLxlK16NQFK90gYEUSSQPDAswJDSG2R/zHn+ +vADI907mW/goEJHeLn8PWGlNlSiR6x+5JJtq+GXCzUzVvJYQSCLGxCSl2x2H+0g7 +NhFI0zPpdzNmO/h+yhzaFb6Rp5U8+ZsnZ3qYjQ/03gw1myTDKJt1YaO9JvArnNYX +hcJQQ8Rt0bHhcrZA16bBOpqZlo5pKCicwI/netgN8QKBgQDcFz7AzdJ26sMSV32V +rwrMgIoggt8qDjO1ARwqW35A1TIge0FoW4M4KpsXQGGfT341uU1esXEcyZ/1L/5X +3ql2gX4DbOYLZLWYzZGR2hq33oi8HkhN98QrEwL9emSH8NqYX3Xxja3PrmCrSYJe +Zbnd9TIm2XkxyMoyXJu6M/QvnwKBgQDc4dzqTbxoGEGa5MuJoGmMwPnqgdG9UM5J +eExVnh7osxc2sOdsiPeRjjQTxs9v2kJwctC359OJoo9yGaaJeSghU4LEWJo1sqnA +fzSCLammYvtVAtniyNv5Mxk/6Uimi4NNDKaAKB+m4K2uSn3U9AmY7KPYMGaSbS9W +XSnobjxm8QKBgC8bPpAvvWs8ZhIn7bY659nLbUT2HeO3dHO6UBf0yzn/J6JyHxbB +93zvCZDZc8uQTRgcmCW7XtVlhjoJUqvl+Wlm39zF0xr/LCsPXKfWAb/2/lcdOCaP +8Emz4QD10EyUTYUtcWYJB/mafhBLRH8F0Nlj4J8WDu2L51MOJTqeYhZLAoGAWffN +icocAbJPlo22sdoa4+/+W5yBF8GAJMDRJtZ+9H1t6SLpQHYRkMIBSETkXUTjZvX9 +Ocs9iIQkNW9pO/mTdO+VBfCo71JUfknR02xR+6m5gYjlws/ZeYlssXGN2/hbhNiw +QOcW7Vv6olFJK6Iy/oz0t6wPO3kpnN3Zogi0paECgYEAwo44M1DdYCtV0snhmYM9 +5u0mPfYt5P2SVLXyUbr+vFTfrTL/WKnXIJgbsnj3Gvf+GIZv9tKcXhSNmEHQCYX4 +X3w9iTPddCHuvZ1fpufi2TyArJh0OkoNtLXJHTKrHjf2N+61AQzFiv5WieJrdE+H +qr32PTUuVGPyO9LyTY4/RL0= +-----END PRIVATE KEY----- +"; diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index a7e7ec80636..7369f9e1db6 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -9,11 +9,13 @@ repository.workspace = true default = [] aws = ["dep:litellm-secrets-aws"] google = ["dep:litellm-secrets-google"] +hashicorp = ["dep:litellm-secrets-hashicorp"] [dependencies] litellm-secrets-types.workspace = true litellm-secrets-aws = { workspace = true, optional = true } litellm-secrets-google = { workspace = true, optional = true } +litellm-secrets-hashicorp = { workspace = true, optional = true } litellm-core-utils.workspace = true base64.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/secrets/README.md b/litellm-rust/crates/secrets/README.md index 183a39e15bb..0d333c7d116 100644 --- a/litellm-rust/crates/secrets/README.md +++ b/litellm-rust/crates/secrets/README.md @@ -9,3 +9,5 @@ Backend failures propagate by default. To allow fallback during a backend failur `get_secret` preserves value types. `get_secret_str` accepts a string default and rejects boolean or JSON values with `Error::TypeMismatch`. `get_secret_bool` accepts a boolean default and converts strings containing `true` or `false`, ignoring surrounding whitespace and ASCII case. Other strings and JSON values produce `Error::TypeMismatch`. Conversion failures never activate fallback or replace a found value with the default Provider payloads remain strings unless explicitly selecting a field from an AWS primary JSON secret. Google caches only successfully decoded string payloads, so reads have identical values and types before and after caching. Confirmed absence and failed reads are not cached. AWS resource-not-found responses and Google HTTP 404 responses indicate absence. Other provider errors remain errors, and successful responses without the required payload are malformed responses rather than missing secrets + +The HashiCorp Vault backend is enabled with the `hashicorp` feature and reads KV v2 values from `HCP_VAULT_*` environment variables. It supports static tokens, AppRole authentication, and TLS certificate authentication diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index 0c6e681b8aa..e5b82580cfe 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -30,4 +30,7 @@ pub enum Error { #[cfg(feature = "google")] #[error(transparent)] Google(#[from] litellm_secrets_google::Error), + #[cfg(feature = "hashicorp")] + #[error(transparent)] + Hashicorp(#[from] litellm_secrets_hashicorp::Error), } diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 943ffdf6158..c035d4cb586 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -13,6 +13,8 @@ pub enum SecretManager { GoogleKms(crate::google::GoogleKms), #[cfg(feature = "google")] GoogleSecretManager(crate::google::GoogleSecretManager), + #[cfg(feature = "hashicorp")] + HashicorpVault(crate::hashicorp::HashicorpVault), } impl SecretManager { @@ -27,6 +29,8 @@ impl SecretManager { Self::GoogleKms(_) => KeyManagementSystem::GoogleKms, #[cfg(feature = "google")] Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager, + #[cfg(feature = "hashicorp")] + Self::HashicorpVault(_) => KeyManagementSystem::HashicorpVault, } } } @@ -78,6 +82,12 @@ pub async fn get_secret_from_manager( .get_secret_from_google_secret_manager(secret_name) .await .map_err(Error::from), + #[cfg(feature = "hashicorp")] + SecretManager::HashicorpVault(client) => client + .async_read_secret(secret_name) + .await + .map(|value| value.map(Secret::String)) + .map_err(Error::from), } } diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index ff2e95f7b2f..6cb391ff844 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -19,3 +19,5 @@ pub use state::{SecretManagerState, secret_manager_would_be_consulted}; pub use litellm_secrets_aws as aws; #[cfg(feature = "google")] pub use litellm_secrets_google as google; +#[cfg(feature = "hashicorp")] +pub use litellm_secrets_hashicorp as hashicorp; diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs index a2cbbd843e1..58941171fd9 100644 --- a/litellm-rust/crates/secrets/tests/handler.rs +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -105,3 +105,131 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites Err(Error::MissingCiphertext) )); } + +#[cfg(feature = "hashicorp")] +#[tokio::test] +async fn hashicorp_handler_resolves_found_missing_and_failed_values() { + use std::sync::Arc; + + use litellm_core_utils::settings::Lookup; + use litellm_secrets::{ + Error, FailurePolicy, KeyManagementSettings, SecretManager, SecretManagerState, + SecretResolver, hashicorp::HashicorpVault, hashicorp::HashicorpVaultConfig, + }; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path}, + }; + + let found_server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/KEY")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"data": {"data": {"key": "remote"}}})), + ) + .mount(&found_server) + .await; + let found_environment: Arc = Arc::new({ + let address = found_server.uri(); + move |name: &str| match name { + "HCP_VAULT_ADDR" => Some(address.clone()), + "HCP_VAULT_TOKEN" => Some("token".into()), + _ => None, + } + }); + let found_config = HashicorpVaultConfig::from_environment(found_environment.as_ref()).unwrap(); + let found_manager = + HashicorpVault::with_client(reqwest::Client::new(), found_config, true).unwrap(); + let found_resolver = SecretResolver::new( + Arc::new(SecretManagerState::new( + SecretManager::HashicorpVault(found_manager), + KeyManagementSettings { + hosted_keys: Some(vec!["KEY".into()]), + ..Default::default() + }, + )), + Arc::new(|_: &str| None), + litellm_secrets::OidcResolver::default(), + ); + assert_eq!( + found_resolver + .get_secret_str("KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + "remote" + ); + + let missing_server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(404)) + .mount(&missing_server) + .await; + let missing_environment: Arc = Arc::new({ + let address = missing_server.uri(); + move |name: &str| match name { + "HCP_VAULT_ADDR" => Some(address.clone()), + "HCP_VAULT_TOKEN" => Some("token".into()), + _ => None, + } + }); + let missing_config = + HashicorpVaultConfig::from_environment(missing_environment.as_ref()).unwrap(); + let missing_manager = + HashicorpVault::with_client(reqwest::Client::new(), missing_config, true).unwrap(); + let missing_state = SecretManagerState::new( + SecretManager::HashicorpVault(missing_manager), + KeyManagementSettings { + hosted_keys: Some(vec!["KEY".into()]), + ..Default::default() + }, + ); + let missing = litellm_secrets::get_secret_from_manager( + missing_state.backend().unwrap(), + "KEY", + missing_state.settings().unwrap(), + &|_: &str| None, + ) + .await + .unwrap(); + assert!(missing.is_none()); + + let failed_server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(500)) + .mount(&failed_server) + .await; + let failed_environment: Arc = Arc::new({ + let address = failed_server.uri(); + move |name: &str| match name { + "HCP_VAULT_ADDR" => Some(address.clone()), + "HCP_VAULT_TOKEN" => Some("token".into()), + _ => None, + } + }); + let failed_config = + HashicorpVaultConfig::from_environment(failed_environment.as_ref()).unwrap(); + let failed_manager = + HashicorpVault::with_client(reqwest::Client::new(), failed_config, true).unwrap(); + let failed_state = SecretManagerState::new( + SecretManager::HashicorpVault(failed_manager), + KeyManagementSettings { + hosted_keys: Some(vec!["KEY".into()]), + ..Default::default() + }, + ); + let failed_resolver = SecretResolver::new( + Arc::new(failed_state), + Arc::new(|_: &str| None), + litellm_secrets::OidcResolver::default(), + ) + .with_failure_policy(FailurePolicy::Propagate); + assert!(matches!( + failed_resolver.get_secret_str("KEY", None).await, + Err(Error::Hashicorp( + litellm_secrets::hashicorp::Error::Status { status: 500 } + )) + )); +} diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index e37a912c7e1..259b2416d7f 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -95,16 +95,16 @@ class HashicorpSecretManager(BaseSecretManager): from litellm.proxy.proxy_server import CommonProxyErrors, premium_user # Vault-specific config - self.vault_addr = os.getenv("HCP_VAULT_ADDR", "http://127.0.0.1:8200") + self.vault_addr = (os.getenv("HCP_VAULT_ADDR") or "http://127.0.0.1:8200").rstrip("/") self.vault_token = os.getenv("HCP_VAULT_TOKEN", "") - self.vault_namespace = os.getenv("HCP_VAULT_NAMESPACE", None) - self.login_namespace_override = os.getenv("HCP_VAULT_LOGIN_NAMESPACE", None) - self.secret_namespace_override = os.getenv("HCP_VAULT_SECRET_NAMESPACE", None) + self.vault_namespace = self._sanitize_path_component(os.getenv("HCP_VAULT_NAMESPACE")) + self.login_namespace_override = self._sanitize_path_component(os.getenv("HCP_VAULT_LOGIN_NAMESPACE")) + self.secret_namespace_override = self._sanitize_path_component(os.getenv("HCP_VAULT_SECRET_NAMESPACE")) # KV engine mount name (default: "secret") # If your KV engine is mounted somewhere other than "secret", set HCP_VAULT_MOUNT_NAME - self.vault_mount_name = os.getenv("HCP_VAULT_MOUNT_NAME", "secret") + self.vault_mount_name = self._sanitize_path_component(os.getenv("HCP_VAULT_MOUNT_NAME")) or "secret" # Optional path prefix for secrets (e.g., "myapp" -> secret/data/myapp/{secret_name}) - self.vault_path_prefix = os.getenv("HCP_VAULT_PATH_PREFIX", None) + self.vault_path_prefix = self._sanitize_path_component(os.getenv("HCP_VAULT_PATH_PREFIX")) # Optional config for TLS cert auth self.tls_cert_path = os.getenv("HCP_VAULT_CLIENT_CERT", "") @@ -114,7 +114,9 @@ class HashicorpSecretManager(BaseSecretManager): # Optional config for AppRole auth self.approle_role_id = os.getenv("HCP_VAULT_APPROLE_ROLE_ID", "") self.approle_secret_id = os.getenv("HCP_VAULT_APPROLE_SECRET_ID", "") - self.approle_mount_path = os.getenv("HCP_VAULT_APPROLE_MOUNT_PATH", "approle") + self.approle_mount_path = self._sanitize_path_component( + os.getenv("HCP_VAULT_APPROLE_MOUNT_PATH") + ) or "approle" self._verify_required_credentials_exist() diff --git a/tests/test_litellm/secret_managers/hashicorp_vault_parity.json b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json new file mode 100644 index 00000000000..5f28d9602a0 --- /dev/null +++ b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json @@ -0,0 +1,97 @@ +[ + { + "name": "defaults", + "env": { + "HCP_VAULT_TOKEN": "token" + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://127.0.0.1:8200/v1/secret/data/OPENAI_API_KEY", + "expected_login_url": null, + "expected_login_namespace": null, + "expected_secret_namespace": null + }, + { + "name": "global_namespace", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_TOKEN": "token", + "HCP_VAULT_NAMESPACE": " admin/ " + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY", + "expected_login_url": null, + "expected_login_namespace": "admin", + "expected_secret_namespace": "admin" + }, + { + "name": "namespace_overrides", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_TOKEN": "token", + "HCP_VAULT_NAMESPACE": "admin", + "HCP_VAULT_LOGIN_NAMESPACE": " /root/ ", + "HCP_VAULT_SECRET_NAMESPACE": " /teams/team-a/ " + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/teams/team-a/secret/data/OPENAI_API_KEY", + "expected_login_url": null, + "expected_login_namespace": "root", + "expected_secret_namespace": "teams/team-a" + }, + { + "name": "custom_mount_and_prefix", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_TOKEN": "token", + "HCP_VAULT_MOUNT_NAME": " /kv-prod/ ", + "HCP_VAULT_PATH_PREFIX": " /virtual-keys/ " + }, + "secret_name": "DB_PASSWORD", + "expected_secret_url": "http://vault.test:8200/v1/kv-prod/data/virtual-keys/DB_PASSWORD", + "expected_login_url": null, + "expected_login_namespace": null, + "expected_secret_namespace": null + }, + { + "name": "approle_custom_mount", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_APPROLE_ROLE_ID": "role-id", + "HCP_VAULT_APPROLE_SECRET_ID": "secret-id", + "HCP_VAULT_APPROLE_MOUNT_PATH": " /custom-approle/ ", + "HCP_VAULT_NAMESPACE": "admin" + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY", + "expected_login_url": "http://vault.test:8200/v1/auth/custom-approle/login", + "expected_login_namespace": "admin", + "expected_secret_namespace": "admin" + }, + { + "name": "tls_cert", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200", + "HCP_VAULT_CLIENT_CERT": "/tmp/client.crt", + "HCP_VAULT_CLIENT_KEY": "/tmp/client.key", + "HCP_VAULT_CERT_ROLE": "vault-role", + "HCP_VAULT_NAMESPACE": "admin" + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY", + "expected_login_url": "http://vault.test:8200/v1/auth/cert/login", + "expected_login_namespace": "admin", + "expected_secret_namespace": "admin" + }, + { + "name": "trailing_address_slash", + "env": { + "HCP_VAULT_ADDR": "http://vault.test:8200///", + "HCP_VAULT_TOKEN": "token" + }, + "secret_name": "OPENAI_API_KEY", + "expected_secret_url": "http://vault.test:8200/v1/secret/data/OPENAI_API_KEY", + "expected_login_url": null, + "expected_login_namespace": null, + "expected_secret_namespace": null + } +] diff --git a/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py b/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py index 1676540e4ec..fc18cb8b8f7 100644 --- a/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py +++ b/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py @@ -1,4 +1,5 @@ import datetime +import json from collections.abc import Mapping from pathlib import Path from typing import Final @@ -18,6 +19,23 @@ LOGIN_RESPONSE: Final = {"auth": {"client_token": "hvs.login-token", "lease_dura SECRET_RESPONSE: Final = {"data": {"data": {"key": "sk-from-vault", "password": "pw-from-vault"}}} NAMESPACE_ENV_VARS: Final = ("HCP_VAULT_NAMESPACE", "HCP_VAULT_LOGIN_NAMESPACE", "HCP_VAULT_SECRET_NAMESPACE") +PARITY_ENV_VARS: Final = ( + "HCP_VAULT_ADDR", + "HCP_VAULT_TOKEN", + "HCP_VAULT_NAMESPACE", + "HCP_VAULT_LOGIN_NAMESPACE", + "HCP_VAULT_SECRET_NAMESPACE", + "HCP_VAULT_MOUNT_NAME", + "HCP_VAULT_PATH_PREFIX", + "HCP_VAULT_APPROLE_ROLE_ID", + "HCP_VAULT_APPROLE_SECRET_ID", + "HCP_VAULT_APPROLE_MOUNT_PATH", + "HCP_VAULT_CLIENT_CERT", + "HCP_VAULT_CLIENT_KEY", + "HCP_VAULT_CERT_ROLE", + "HCP_VAULT_REFRESH_INTERVAL", + "SECRET_MANAGER_REFRESH_INTERVAL", +) def _build_manager(monkeypatch: pytest.MonkeyPatch, env: Mapping[str, str]) -> HashicorpSecretManager: @@ -236,3 +254,35 @@ def test_tls_login_uses_login_namespace(monkeypatch: pytest.MonkeyPatch, tmp_pat assert manager._auth_via_tls_cert() == "hvs.login-token" assert login_route.calls.last.request.headers["X-Vault-Namespace"] == "root" + + +with Path(__file__).with_name("hashicorp_vault_parity.json").open() as parity_file: + PARITY_CASES: Final = json.load(parity_file) + + +@pytest.mark.parametrize("case", PARITY_CASES, ids=lambda case: case["name"]) +def test_configuration_matches_native_parity_fixture( + monkeypatch: pytest.MonkeyPatch, case: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + for name in PARITY_ENV_VARS: + monkeypatch.delenv(name, raising=False) + for name, value in case["env"].items(): + monkeypatch.setenv(name, value) + + manager: Final = HashicorpSecretManager() + env: Final = case["env"] + expected_login_url: Final = case["expected_login_url"] + if env.get("HCP_VAULT_APPROLE_ROLE_ID") and env.get("HCP_VAULT_APPROLE_SECRET_ID"): + login_url: str | None = ( + f"{manager.vault_addr}/v1/auth/{manager.approle_mount_path}/login" + ) + elif env.get("HCP_VAULT_CLIENT_CERT") and env.get("HCP_VAULT_CLIENT_KEY"): + login_url = f"{manager.vault_addr}/v1/auth/cert/login" + else: + login_url = None + + assert manager.get_url(case["secret_name"]) == case["expected_secret_url"] + assert manager.vault_login_namespace == case["expected_login_namespace"] + assert manager.vault_secret_namespace == case["expected_secret_namespace"] + assert login_url == expected_login_url From ca31149040170297c1a00044d6bb9650862dd895 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:28:27 +0000 Subject: [PATCH 04/14] feat(rust): add native GCS object-store cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 17 + litellm-rust/Cargo.toml | 2 + litellm-rust/crates/auth-gcp/src/lib.rs | 8 + litellm-rust/crates/cache-gcs/Cargo.toml | 21 ++ litellm-rust/crates/cache-gcs/src/cache.rs | 285 +++++++++++++++++ litellm-rust/crates/cache-gcs/src/lib.rs | 5 + litellm-rust/crates/cache-gcs/src/token.rs | 44 +++ litellm-rust/crates/cache-gcs/tests/cache.rs | 290 ++++++++++++++++++ litellm-rust/crates/cache/src/error.rs | 2 + litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/cache/config.rs | 143 ++++++++- .../crates/python-bridge/src/cache/facade.rs | 4 + .../crates/python-bridge/src/cache/handle.rs | 27 ++ .../crates/python-bridge/src/cache/mod.rs | 3 +- .../crates/python-bridge/src/cache/native.rs | 35 +++ 15 files changed, 876 insertions(+), 11 deletions(-) create mode 100644 litellm-rust/crates/cache-gcs/Cargo.toml create mode 100644 litellm-rust/crates/cache-gcs/src/cache.rs create mode 100644 litellm-rust/crates/cache-gcs/src/lib.rs create mode 100644 litellm-rust/crates/cache-gcs/src/token.rs create mode 100644 litellm-rust/crates/cache-gcs/tests/cache.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..99a6133b4f1 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2464,6 +2464,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-gcs" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-auth-gcp", + "litellm-auth-types", + "litellm-cache", + "percent-encoding", + "reqwest 0.12.28", + "serde_json", + "tokio", + "url", + "wiremock", +] + [[package]] name = "litellm-cache-memory" version = "0.1.0" @@ -2666,6 +2682,7 @@ dependencies = [ "litellm-auth", "litellm-auth-gcp", "litellm-cache", + "litellm-cache-gcs", "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-response", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..80425d567c0 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -29,6 +29,7 @@ litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-redis = { path = "crates/cache-redis" } +litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } @@ -66,6 +67,7 @@ base64 = "0.22" moka = { version = "0.12.16", features = ["future"] } strum = { version = "0.28.0", features = ["derive"] } url = "2.5.8" +percent-encoding = "2.3" webpki-roots = "1" time = { version = "0.3.53", features = ["parsing"] } criterion = "0.8.2" diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 8aeddae9efc..534d85acdb0 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -128,6 +128,14 @@ impl VertexAuth { } } + pub async fn access_token( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + self.load_provider(config, env_lookup).await?.token().await + } + pub async fn validate_environment( &self, headers: Vec<(String, String)>, diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml new file mode 100644 index 00000000000..3890e60843c --- /dev/null +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-cache-gcs" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-auth-gcp.workspace = true +litellm-auth-types.workspace = true +litellm-cache.workspace = true +percent-encoding.workspace = true +reqwest.workspace = true +tokio.workspace = true +url.workspace = true + +[dev-dependencies] +serde_json.workspace = true +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs new file mode 100644 index 00000000000..1097f3252b1 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -0,0 +1,285 @@ +use std::{future::Future, sync::Arc, time::Duration}; + +use futures_util::future::try_join_all; +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, + FlushCache, +}; +use percent_encoding::{AsciiSet, CONTROLS, percent_encode}; +use reqwest::Client; + +use crate::{GcpTokenSource, TokenSource}; + +pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; + +const OBJECT_NAME_ENCODE_SET: &AsciiSet = &CONTROLS + .add(b' ') + .add(b'!') + .add(b'"') + .add(b'#') + .add(b'$') + .add(b'%') + .add(b'&') + .add(b'\'') + .add(b'(') + .add(b')') + .add(b'*') + .add(b'+') + .add(b',') + .add(b'/') + .add(b':') + .add(b';') + .add(b'<') + .add(b'=') + .add(b'>') + .add(b'?') + .add(b'@') + .add(b'[') + .add(b'\\') + .add(b']') + .add(b'^') + .add(b'`') + .add(b'{') + .add(b'|') + .add(b'}'); + +pub fn key_prefix(gcs_path: Option<&str>) -> String { + match gcs_path { + Some(path) if !path.is_empty() => format!("{}/", path.trim_end_matches('/')), + _ => String::new(), + } +} + +#[derive(Clone, Debug)] +pub struct GcsConfig { + pub bucket_name: String, + pub gcs_path: Option, + pub path_service_account: Option, + pub endpoint: String, +} + +impl GcsConfig { + pub fn new(bucket_name: impl Into) -> Self { + Self { + bucket_name: bucket_name.into(), + gcs_path: None, + path_service_account: None, + endpoint: DEFAULT_ENDPOINT.to_string(), + } + } +} + +pub struct GcsCache { + config: GcsConfig, + key_prefix: String, + client: Client, + token: Arc, + codec: S, +} + +impl GcsCache { + pub fn new(config: GcsConfig, codec: S) -> Result { + let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone())); + Self::with_token_source(config, codec, token) + } + + pub fn with_token_source( + config: GcsConfig, + codec: S, + token: Arc, + ) -> Result { + let client = Client::builder().build().map_err(|_| Error::Unavailable)?; + let key_prefix = key_prefix(config.gcs_path.as_deref()); + Ok(Self { + config, + key_prefix, + client, + token, + codec, + }) + } + + pub fn bucket_name(&self) -> &str { + &self.config.bucket_name + } + + pub fn key_prefix(&self) -> &str { + &self.key_prefix + } + + pub fn path_service_account(&self) -> Option<&str> { + self.config.path_service_account.as_deref() + } + + pub fn object_name(&self, key: &str) -> String { + format!("{}{}", self.key_prefix, key) + } + + fn encoded_object_name(&self, key: &str) -> String { + percent_encode(self.object_name(key).as_bytes(), OBJECT_NAME_ENCODE_SET).to_string() + } + + fn endpoint(&self, path: &str) -> String { + format!("{}{}", self.config.endpoint.trim_end_matches('/'), path) + } + + async fn async_set(&self, key: &str, value: S::Value) -> Result<(), Error> { + let token = self.token.bearer_token().await?; + let payload = self.codec.encode(&value)?; + let url = self.endpoint(&format!( + "/upload/storage/v1/b/{}/o?uploadType=media&name={}", + self.config.bucket_name, + self.encoded_object_name(key) + )); + let response = self + .client + .post(url) + .bearer_auth(token) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .body(payload) + .send() + .await + .map_err(|_| Error::Unavailable)?; + if !response.status().is_success() { + return Err(Error::Unavailable); + } + Ok(()) + } + + async fn async_get(&self, key: &str) -> Result, Error> { + let token = self.token.bearer_token().await?; + let url = self.endpoint(&format!( + "/storage/v1/b/{}/o/{}?alt=media", + self.config.bucket_name, + self.encoded_object_name(key) + )); + let response = self + .client + .get(url) + .bearer_auth(token) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .send() + .await + .map_err(|_| Error::Unavailable)?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + if !response.status().is_success() { + return Err(Error::Unavailable); + } + let body = response.bytes().await.map_err(|_| Error::Unavailable)?; + self.codec + .decode(&body) + .map(Some) + .map_err(|_| Error::InvalidEntry) + } + + fn run_sync(future: F) -> Result + where + F: Future> + Send, + T: Send, + { + let run = || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|_| Error::Unavailable) + .and_then(|runtime| runtime.block_on(future)) + }; + if let Ok(handle) = tokio::runtime::Handle::try_current() { + if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread { + return tokio::task::block_in_place(run); + } + return std::thread::scope(|scope| { + scope + .spawn(run) + .join() + .map_err(|_| Error::Unavailable) + .and_then(|result| result) + }); + } + run() + } +} + +impl BaseCache for GcsCache { + type Value = S::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + None + } + + fn set_cache(&self, key: &str, value: Self::Value, _: &Self::Context) -> Result<(), Error> { + Self::run_sync(self.async_set(key, value)) + } + + fn get_cache(&self, key: &str, _: &Self::Context) -> Result, Error> { + Self::run_sync(self.async_get(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + _: Self::Context, + ) -> Result<(), Error> { + self.async_set(key, value).await + } + + async fn async_get_cache( + &self, + key: &str, + _: &Self::Context, + ) -> Result, Error> { + self.async_get(key).await + } + + async fn async_set_cache_pipeline( + &self, + entries: Vec<(String, Self::Value)>, + context: Self::Context, + ) -> Result<(), Error> { + try_join_all(entries.into_iter().map(|(key, value)| { + let context = context.clone(); + async move { self.async_set_cache(&key, value, context).await } + })) + .await + .map(|_| ()) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} + +impl BatchCache for GcsCache { + async fn async_batch_get_cache( + &self, + keys: Vec, + context: Self::Context, + ) -> Result>, Error> { + try_join_all(keys.into_iter().map(|key| { + let context = context.clone(); + async move { + match self.async_get_cache(&key, &context).await { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + } + } + })) + .await + } +} + +impl FlushCache for GcsCache { + fn flush_cache(&self) -> Result<(), Error> { + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-gcs/src/lib.rs b/litellm-rust/crates/cache-gcs/src/lib.rs new file mode 100644 index 00000000000..cbb61cf0685 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/lib.rs @@ -0,0 +1,5 @@ +mod cache; +mod token; + +pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix}; +pub use token::{GcpTokenSource, StaticTokenSource, TokenSource}; diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs new file mode 100644 index 00000000000..adb601c276c --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -0,0 +1,44 @@ +use std::{future::Future, pin::Pin}; + +use litellm_auth_gcp::{VertexAuth, VertexConfig}; +use litellm_auth_types::{InputSource, SecretValue, Sourced}; +use litellm_cache::Error; + +pub trait TokenSource: Send + Sync + 'static { + fn bearer_token(&self) -> Pin> + Send + '_>>; +} + +pub struct GcpTokenSource { + auth: VertexAuth, + config: VertexConfig, +} + +impl GcpTokenSource { + pub fn new(path_service_account: Option) -> Self { + let credentials = path_service_account + .map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment)); + Self { + auth: VertexAuth::default(), + config: VertexConfig::new(credentials, None, None), + } + } +} + +impl TokenSource for GcpTokenSource { + fn bearer_token(&self) -> Pin> + Send + '_>> { + Box::pin(async move { + self.auth + .access_token(&self.config, &|name| std::env::var(name).ok()) + .await + .map_err(|_| Error::Unavailable) + }) + } +} + +pub struct StaticTokenSource(pub String); + +impl TokenSource for StaticTokenSource { + fn bearer_token(&self) -> Pin> + Send + '_>> { + Box::pin(async move { Ok(self.0.clone()) }) + } +} diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs new file mode 100644 index 00000000000..79e4ed100f8 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -0,0 +1,290 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheContext, Error, ExactCacheContext, FlushCache, + JsonCodec, +}; +use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource, key_prefix}; +use serde_json::json; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_bytes, header, method, path, query_param}, +}; + +fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig { + GcsConfig { + bucket_name: "bucket".into(), + gcs_path: gcs_path.map(str::to_string), + path_service_account: None, + endpoint: server.uri(), + } +} + +fn cache(server: &MockServer, gcs_path: Option<&str>) -> GcsCache> { + GcsCache::with_token_source( + config(server, gcs_path), + JsonCodec::new(), + Arc::new(StaticTokenSource("tok".into())), + ) + .unwrap() +} + +#[tokio::test] +async fn set_writes_encoded_object_and_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .and(header("authorization", "Bearer tok")) + .and(header("content-type", "application/json")) + .and(body_bytes(br#"{"value":"entry"}"#)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + cache(&server, Some("cache/")) + .set_cache( + "team:a b/c", + json!({"value": "entry"}), + &ExactCacheContext::default(), + ) + .unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].url.query(), + Some("uploadType=media&name=cache%2Fteam%3Aa%20b%2Fc") + ); +} + +#[tokio::test] +async fn get_maps_statuses_and_decode_failures() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/hit")) + .and(query_param("alt", "media")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/missing")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/server-error")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/invalid")) + .respond_with(ResponseTemplate::new(200).set_body_string("not json")) + .mount(&server) + .await; + + let cache = cache(&server, None); + assert_eq!( + cache + .get_cache("hit", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"value": "entry"})) + ); + assert_eq!( + cache + .get_cache("missing", &ExactCacheContext::default()) + .unwrap(), + None + ); + assert_eq!( + cache + .get_cache("server-error", &ExactCacheContext::default()) + .unwrap_err(), + Error::Unavailable + ); + assert_eq!( + cache + .get_cache("invalid", &ExactCacheContext::default()) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn key_prefix_normalizes_paths() { + assert_eq!(key_prefix(None), ""); + assert_eq!(key_prefix(Some("a/b/")), "a/b/"); + assert_eq!(key_prefix(Some("a/b")), "a/b/"); + assert_eq!(key_prefix(Some("")), ""); +} + +#[tokio::test] +async fn ignores_ttl_and_writes_pipeline_concurrently() { + let server = MockServer::start().await; + for key in ["one", "two", "three"] { + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .and(query_param("name", key)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + } + let cache = cache(&server, None); + assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5)))), + None + ); + cache + .async_set_cache_pipeline( + vec![ + ("one".into(), json!({"key": "one"})), + ("two".into(), json!({"key": "two"})), + ("three".into(), json!({"key": "three"})), + ], + ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5))), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn async_batch_get_preserves_hits_misses_and_invalid_entries() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/hit")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/missing")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/invalid")) + .respond_with(ResponseTemplate::new(200).set_body_string("not json")) + .mount(&server) + .await; + + assert_eq!( + cache(&server, None) + .async_batch_get_cache( + vec!["hit".into(), "missing".into(), "invalid".into()], + ExactCacheContext::default(), + ) + .await + .unwrap(), + vec![ + BatchEntry::Hit(json!({"value": "entry"})), + BatchEntry::Miss, + BatchEntry::Invalid, + ] + ); +} + +#[tokio::test] +async fn lifecycle_operations_are_noops_and_connection_test_is_unsupported() { + let server = MockServer::start().await; + let cache = cache(&server, None); + assert_eq!(cache.flush_cache(), Ok(())); + assert_eq!(cache.disconnect().await, Ok(())); + assert_eq!( + cache.test_connection().await, + Err(Error::UnsupportedOperation) + ); +} + +#[test] +fn sync_operations_work_without_an_active_runtime() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + let server = runtime.block_on(MockServer::start()); + runtime.block_on( + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server), + ); + runtime.block_on( + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server), + ); + let cache = cache(&server, None); + cache + .set_cache( + "key", + json!({"value": "entry"}), + &ExactCacheContext::default(), + ) + .unwrap(); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"value": "entry"})) + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn sync_operations_work_inside_a_multi_thread_runtime() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server) + .await; + let cache = cache(&server, None); + cache + .set_cache( + "key", + json!({"value": "entry"}), + &ExactCacheContext::default(), + ) + .unwrap(); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"value": "entry"})) + ); +} + +struct FailingTokenSource; + +impl TokenSource for FailingTokenSource { + fn bearer_token( + &self, + ) -> std::pin::Pin> + Send + '_>> + { + Box::pin(async { Err(Error::Unavailable) }) + } +} + +#[tokio::test] +async fn token_source_failure_skips_http() { + let server = MockServer::start().await; + let cache = GcsCache::with_token_source( + config(&server, None), + JsonCodec::::new(), + Arc::new(FailingTokenSource), + ) + .unwrap(); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap_err(), + Error::Unavailable + ); + assert_eq!(server.received_requests().await.unwrap().len(), 0); +} diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..1a381d0afd8 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,4 +6,6 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, + #[error("operation is not supported by this cache")] + UnsupportedOperation, } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..0f060c801a5 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -23,6 +23,7 @@ bytes.workspace = true litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true +litellm-cache-gcs.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 0e7d6aee11d..66f493dd760 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -73,9 +73,17 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } +#[derive(Debug, PartialEq)] +pub(super) struct GcsCacheConfig { + pub(super) bucket_name: String, + pub(super) key_prefix: String, + pub(super) path_service_account: Option, +} + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + Gcs(GcsCacheConfig), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -90,6 +98,7 @@ pub(super) enum UnsupportedCacheConfig { RedisCredentials, RedisConnection, RedisOption, + GcsBucket, } impl UnsupportedCacheConfig { @@ -100,6 +109,7 @@ impl UnsupportedCacheConfig { Self::RedisCredentials => "native Redis credentials require Python", Self::RedisConnection => "native Redis connection type is not implemented", Self::RedisOption => "native Redis configuration requires Python", + Self::GcsBucket => "native GCS cache requires a configured bucket name", } } } @@ -142,14 +152,20 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::Gcs) => match project_gcs(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::Gcs(backend), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some( CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::S3 | CacheType::Disk | CacheType::QdrantSemantic - | CacheType::AzureBlob - | CacheType::Gcs, + | CacheType::AzureBlob, ) | None => Ok(CacheConfigProjection::Unsupported( UnsupportedCacheConfig::Backend, @@ -158,12 +174,12 @@ impl NativeCacheConfig { } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - if service.default_ttl() - != Some(match &self.backend { - CacheBackendConfig::Memory(config) => config.default_ttl, - CacheBackendConfig::Redis(config) => config.default_ttl, - }) - { + let expected = match &self.backend { + CacheBackendConfig::Memory(config) => Some(config.default_ttl), + CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::Gcs(_) => None, + }; + if service.default_ttl() != expected { return Some("facade and native backend default TTLs must match"); } match &self.backend { @@ -185,6 +201,31 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::Gcs(config) + if service + .gcs_backend() + .is_none_or(|backend| backend.bucket_name() != config.bucket_name) => + { + Some("facade and native backend buckets must match") + } + CacheBackendConfig::Gcs(config) + if service + .gcs_backend() + .is_none_or(|backend| backend.key_prefix() != config.key_prefix) => + { + Some("facade and native backend key prefixes must match") + } + CacheBackendConfig::Gcs(config) + if service.gcs_backend().is_none_or(|backend| { + backend.path_service_account() != config.path_service_account.as_deref() + }) => + { + Some("facade and native backend credentials must match") + } + CacheBackendConfig::Gcs(_) => None, } } } @@ -201,6 +242,23 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { }) } +#[inline(never)] +fn project_gcs( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let bucket_name = match backend.getattr("bucket_name")?.extract::>() { + Ok(Some(bucket_name)) if !bucket_name.is_empty() => bucket_name, + _ => return Ok(Err(UnsupportedCacheConfig::GcsBucket)), + }; + Ok(Ok(GcsCacheConfig { + bucket_name, + key_prefix: backend.getattr("key_prefix")?.extract::()?, + path_service_account: backend + .getattr("path_service_account")? + .extract::>()?, + })) +} + #[inline(never)] fn project_redis( backend: &Bound<'_, PyAny>, @@ -468,8 +526,8 @@ mod tests { use pyo3::{prelude::*, types::PyDict}; use super::{ - CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig, - RedisProtocol, + CacheBackendConfig, CacheConfigProjection, CertificateRequirement, GcsCacheConfig, + NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig, }; use crate::cache::native::NativeResponseCache; @@ -531,6 +589,71 @@ mod tests { }); } + #[test] + fn projects_gcs_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(bucket_name='bucket', key_prefix='cache/', path_service_account='credentials.json')\n\ + facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("GCS cache should be supported"); + }; + let CacheBackendConfig::Gcs(gcs) = config.backend else { + panic!("expected GCS configuration"); + }; + assert_eq!( + gcs, + GcsCacheConfig { + bucket_name: "bucket".into(), + key_prefix: "cache/".into(), + path_service_account: Some("credentials.json".into()), + } + ); + let matching = NativeResponseCache::gcs( + litellm_cache_gcs::GcsConfig { + bucket_name: "bucket".into(), + gcs_path: Some("cache/".into()), + path_service_account: Some("credentials.json".into()), + endpoint: litellm_cache_gcs::DEFAULT_ENDPOINT.into(), + }, + Some("token".into()), + ) + .unwrap(); + let matching_config = NativeCacheConfig { + policy: config.policy, + backend: CacheBackendConfig::Gcs(gcs), + }; + assert_eq!(matching_config.service_mismatch(&matching), None); + }); + } + + #[test] + fn rejects_gcs_without_a_bucket_name() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(bucket_name=None, key_prefix='', path_service_account=None)\n\ + facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("GCS cache without a bucket should be unsupported"); + }; + assert!(matches!(&reason, UnsupportedCacheConfig::GcsBucket)); + assert_eq!( + reason.message(), + "native GCS cache requires a configured bucket name" + ); + }); + } + #[test] fn projects_resolved_redis_tls_configuration() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f2f86c14b37..6aa94754e2d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -192,6 +192,7 @@ impl FacadeGuard { let (module, name, cache_kind) = match kind { "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + "gcs" => ("litellm.caching.gcs_cache", "GCSCache", "gcs"), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -235,6 +236,9 @@ impl FacadeGuard { "max_size_per_item", "redis_kwargs", "redis_flush_size", + "bucket_name", + "key_prefix", + "path_service_account", ], )?, redis_pool: (kind == "redis") diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8251b3df06c..87d9300b08b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,6 +1,8 @@ use litellm_host_python::release_gil; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; + use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; #[pyclass(frozen, name = "_CacheTestHandle")] @@ -51,6 +53,31 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (bucket_name, *, gcs_path=None, path_service_account=None, endpoint=None, token=None))] + fn gcs( + py: Python<'_>, + bucket_name: String, + gcs_path: Option, + path_service_account: Option, + endpoint: Option, + token: Option, + ) -> PyResult { + let config = GcsConfig { + bucket_name, + gcs_path, + path_service_account, + endpoint: endpoint.unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()), + }; + let service = release_gil(py, move || NativeResponseCache::gcs(config, token)) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index aec08610f6e..ac4494150d9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -10,7 +10,7 @@ mod resolver; use litellm_cache::Error; use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, }; @@ -21,6 +21,7 @@ pub(crate) use self::{ fn cache_error(error: Error) -> PyErr { match error { Error::InvalidEntry => PyValueError::new_err(error.to_string()), + Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()), _ => PyRuntimeError::new_err(error.to_string()), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a9475429e45..9a030ac0434 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,6 +1,7 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ @@ -15,6 +16,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + Gcs(Arc>>), } impl NativeResponseCache { @@ -43,6 +45,18 @@ impl NativeResponseCache { buffer: None, }) } + + pub fn gcs(config: GcsConfig, token: Option) -> Result { + let backend = match token { + Some(token) => GcsCache::with_token_source( + config, + ResponseCacheCodec, + Arc::new(StaticTokenSource(token)), + )?, + None => GcsCache::new(config, ResponseCacheCodec)?, + }; + Ok(Self::Gcs(Arc::new(ResponseCache::new(Arc::new(backend))))) + } } impl NativeResponseCache { @@ -50,6 +64,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::Gcs(_) => "gcs", } } @@ -57,6 +72,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::Gcs(cache) => cache.default_ttl(), } } @@ -64,6 +80,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), + Self::Gcs(_) => None, } } @@ -71,6 +88,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), Self::Redis { .. } => None, + Self::Gcs(_) => None, } } @@ -78,6 +96,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), Self::Redis { .. } => None, + Self::Gcs(_) => None, } } @@ -99,6 +118,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup(request, now), Self::Redis { cache, .. } => cache.lookup(request, now), + Self::Gcs(cache) => cache.lookup(request, now), } } @@ -111,6 +131,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.store(request, response, now), Self::Redis { cache, .. } => cache.store(request, response, now), + Self::Gcs(cache) => cache.store(request, response, now), } } @@ -122,6 +143,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup_batch(requests, now), Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::Gcs(cache) => cache.lookup_batch(requests, now), } } @@ -133,6 +155,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup(request, now).await, Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::Gcs(cache) => cache.async_lookup(request, now).await, } } @@ -152,6 +175,7 @@ impl NativeResponseCache { cache, buffer: Some(buffer), } => buffer.async_store(cache, request, response, now).await, + Self::Gcs(cache) => cache.async_store(request, response, now).await, } } @@ -163,6 +187,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await, } } @@ -174,6 +199,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_store_batch(entries, now).await, Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + Self::Gcs(cache) => cache.async_store_batch(entries, now).await, } } @@ -186,6 +212,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::Gcs(cache) => cache.async_flush().await, } } @@ -193,6 +220,14 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::Gcs(cache) => cache.test_connection().await, + } + } + + pub fn gcs_backend(&self) -> Option<&GcsCache> { + match self { + Self::Gcs(cache) => Some(cache.backend()), + _ => None, } } } From beba2576be4b015574dccb3f916af5dfaf5d189f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:29:06 +0000 Subject: [PATCH 05/14] fix(rust): align vault namespace handling with python and drop unused settings hook Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/secrets-hashicorp/src/config.rs | 10 +----- .../secrets-hashicorp/src/secret_manager.rs | 31 ++++++++--------- .../secrets-hashicorp/tests/secret_manager.rs | 33 +++++++++++++++++++ .../hashicorp_secret_manager.py | 16 ++++----- .../hashicorp_vault_parity.json | 20 +++-------- 5 files changed, 58 insertions(+), 52 deletions(-) diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs index f9491f71afb..d32491199c4 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/config.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -1,7 +1,6 @@ use std::{path::PathBuf, time::Duration}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::KeyManagementSettings; use litellm_secrets_types::SecretValue; use crate::Error; @@ -117,13 +116,6 @@ impl HashicorpVaultConfig { }) } - pub fn from_settings( - _settings: &KeyManagementSettings, - environment: &dyn Lookup, - ) -> Result { - Self::from_environment(environment) - } - pub fn login_namespace(&self) -> Option<&str> { self.login_namespace .as_deref() @@ -163,7 +155,7 @@ fn refresh_interval(environment: &dyn Lookup) -> Result { }; let seconds: i64 = value.parse().map_err(|_| Error::RefreshInterval)?; if seconds < 0 { - return Ok(Duration::from_nanos(1)); + return Err(Error::RefreshInterval); } Ok(Duration::from_secs(seconds as u64)) } diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index 6b887289443..e4ca8c387c3 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -120,14 +120,12 @@ impl HashicorpVault { return Ok(Some(value)); } let token: SecretValue = self.vault_token().await?; - let mut request = self + let response: reqwest::Response = self .client .get(&url) - .header("X-Vault-Token", token.expose()); - if let Some(namespace) = self.config.secret_namespace() { - request = request.header("X-Vault-Namespace", namespace); - } - let response: reqwest::Response = request.send().await?; + .header("X-Vault-Token", token.expose()) + .send() + .await?; if response.status() == reqwest::StatusCode::NOT_FOUND { return Ok(None); } @@ -166,14 +164,13 @@ impl HashicorpVault { None => json!({"key": value.expose()}), }; let token: SecretValue = self.vault_token().await?; - let mut request = self + let response: reqwest::Response = self .client .post(&url) - .header("X-Vault-Token", token.expose()); - if let Some(namespace) = self.config.secret_namespace() { - request = request.header("X-Vault-Namespace", namespace); - } - let response: reqwest::Response = request.json(&json!({"data": data})).send().await?; + .header("X-Vault-Token", token.expose()) + .json(&json!({"data": data})) + .send() + .await?; if !response.status().is_success() { return Err(Error::Status { status: response.status().as_u16(), @@ -190,14 +187,12 @@ impl HashicorpVault { pub async fn async_delete_secret(&self, secret_name: &str) -> Result<(), Error> { let url: String = self.secret_url(secret_name)?; let token: SecretValue = self.vault_token().await?; - let mut request = self + let response: reqwest::Response = self .client .delete(&url) - .header("X-Vault-Token", token.expose()); - if let Some(namespace) = self.config.secret_namespace() { - request = request.header("X-Vault-Namespace", namespace); - } - let response: reqwest::Response = request.send().await?; + .header("X-Vault-Token", token.expose()) + .send() + .await?; if !response.status().is_success() { return Err(Error::Status { status: response.status().as_u16(), diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs index 1cf13a79c77..b7a81f49ad4 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -87,6 +87,39 @@ async fn namespace_mount_and_prefix_are_sanitized_in_the_url() { assert!(manager.async_read_secret("name").await.unwrap().is_some()); } +#[test] +fn trailing_address_slashes_are_removed() { + let environment: Arc = Arc::new(|name: &str| match name { + "HCP_VAULT_ADDR" => Some("http://vault.test:8200///".to_owned()), + "HCP_VAULT_TOKEN" => Some("token".to_owned()), + _ => None, + }); + let config: HashicorpVaultConfig = + HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap(); + let manager: HashicorpVault = + HashicorpVault::with_client(reqwest::Client::new(), config, true).unwrap(); + + assert_eq!( + manager.secret_url("name").unwrap(), + "http://vault.test:8200/v1/secret/data/name" + ); +} + +#[rstest::rstest] +#[case("-1")] +#[case("not-a-number")] +fn invalid_refresh_intervals_are_rejected(#[case] value: &str) { + let environment: Arc = Arc::new(move |name: &str| match name { + "HCP_VAULT_REFRESH_INTERVAL" => Some(value.to_owned()), + _ => None, + }); + + assert!(matches!( + HashicorpVaultConfig::from_environment(environment.as_ref()), + Err(Error::RefreshInterval) + )); +} + #[tokio::test] async fn approle_login_uses_namespace_and_reuses_the_token() { let server: MockServer = MockServer::start().await; diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index 259b2416d7f..e37a912c7e1 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -95,16 +95,16 @@ class HashicorpSecretManager(BaseSecretManager): from litellm.proxy.proxy_server import CommonProxyErrors, premium_user # Vault-specific config - self.vault_addr = (os.getenv("HCP_VAULT_ADDR") or "http://127.0.0.1:8200").rstrip("/") + self.vault_addr = os.getenv("HCP_VAULT_ADDR", "http://127.0.0.1:8200") self.vault_token = os.getenv("HCP_VAULT_TOKEN", "") - self.vault_namespace = self._sanitize_path_component(os.getenv("HCP_VAULT_NAMESPACE")) - self.login_namespace_override = self._sanitize_path_component(os.getenv("HCP_VAULT_LOGIN_NAMESPACE")) - self.secret_namespace_override = self._sanitize_path_component(os.getenv("HCP_VAULT_SECRET_NAMESPACE")) + self.vault_namespace = os.getenv("HCP_VAULT_NAMESPACE", None) + self.login_namespace_override = os.getenv("HCP_VAULT_LOGIN_NAMESPACE", None) + self.secret_namespace_override = os.getenv("HCP_VAULT_SECRET_NAMESPACE", None) # KV engine mount name (default: "secret") # If your KV engine is mounted somewhere other than "secret", set HCP_VAULT_MOUNT_NAME - self.vault_mount_name = self._sanitize_path_component(os.getenv("HCP_VAULT_MOUNT_NAME")) or "secret" + self.vault_mount_name = os.getenv("HCP_VAULT_MOUNT_NAME", "secret") # Optional path prefix for secrets (e.g., "myapp" -> secret/data/myapp/{secret_name}) - self.vault_path_prefix = self._sanitize_path_component(os.getenv("HCP_VAULT_PATH_PREFIX")) + self.vault_path_prefix = os.getenv("HCP_VAULT_PATH_PREFIX", None) # Optional config for TLS cert auth self.tls_cert_path = os.getenv("HCP_VAULT_CLIENT_CERT", "") @@ -114,9 +114,7 @@ class HashicorpSecretManager(BaseSecretManager): # Optional config for AppRole auth self.approle_role_id = os.getenv("HCP_VAULT_APPROLE_ROLE_ID", "") self.approle_secret_id = os.getenv("HCP_VAULT_APPROLE_SECRET_ID", "") - self.approle_mount_path = self._sanitize_path_component( - os.getenv("HCP_VAULT_APPROLE_MOUNT_PATH") - ) or "approle" + self.approle_mount_path = os.getenv("HCP_VAULT_APPROLE_MOUNT_PATH", "approle") self._verify_required_credentials_exist() diff --git a/tests/test_litellm/secret_managers/hashicorp_vault_parity.json b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json index 5f28d9602a0..f1faefd48e8 100644 --- a/tests/test_litellm/secret_managers/hashicorp_vault_parity.json +++ b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json @@ -15,7 +15,7 @@ "env": { "HCP_VAULT_ADDR": "http://vault.test:8200", "HCP_VAULT_TOKEN": "token", - "HCP_VAULT_NAMESPACE": " admin/ " + "HCP_VAULT_NAMESPACE": "admin" }, "secret_name": "OPENAI_API_KEY", "expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY", @@ -29,8 +29,8 @@ "HCP_VAULT_ADDR": "http://vault.test:8200", "HCP_VAULT_TOKEN": "token", "HCP_VAULT_NAMESPACE": "admin", - "HCP_VAULT_LOGIN_NAMESPACE": " /root/ ", - "HCP_VAULT_SECRET_NAMESPACE": " /teams/team-a/ " + "HCP_VAULT_LOGIN_NAMESPACE": "root", + "HCP_VAULT_SECRET_NAMESPACE": "teams/team-a" }, "secret_name": "OPENAI_API_KEY", "expected_secret_url": "http://vault.test:8200/v1/teams/team-a/secret/data/OPENAI_API_KEY", @@ -58,7 +58,7 @@ "HCP_VAULT_ADDR": "http://vault.test:8200", "HCP_VAULT_APPROLE_ROLE_ID": "role-id", "HCP_VAULT_APPROLE_SECRET_ID": "secret-id", - "HCP_VAULT_APPROLE_MOUNT_PATH": " /custom-approle/ ", + "HCP_VAULT_APPROLE_MOUNT_PATH": "custom-approle", "HCP_VAULT_NAMESPACE": "admin" }, "secret_name": "OPENAI_API_KEY", @@ -81,17 +81,5 @@ "expected_login_url": "http://vault.test:8200/v1/auth/cert/login", "expected_login_namespace": "admin", "expected_secret_namespace": "admin" - }, - { - "name": "trailing_address_slash", - "env": { - "HCP_VAULT_ADDR": "http://vault.test:8200///", - "HCP_VAULT_TOKEN": "token" - }, - "secret_name": "OPENAI_API_KEY", - "expected_secret_url": "http://vault.test:8200/v1/secret/data/OPENAI_API_KEY", - "expected_login_url": null, - "expected_login_namespace": null, - "expected_secret_namespace": null } ] From a6d1497932e257a809fdd4647d85a70039376816 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:38:20 +0000 Subject: [PATCH 06/14] test(rust): add GCS native cache parity fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 - litellm-rust/crates/cache-gcs/Cargo.toml | 1 - litellm-rust/crates/cache-gcs/src/cache.rs | 37 +-- litellm-rust/crates/cache-gcs/tests/cache.rs | 34 +++ tests/test_litellm_rust/support/fake_gcs.py | 152 +++++++++++++ tests/test_litellm_rust/test_cache.py | 228 +++++++++++++++++++ 6 files changed, 420 insertions(+), 33 deletions(-) create mode 100644 tests/test_litellm_rust/support/fake_gcs.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 99a6133b4f1..ae28553801d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2476,7 +2476,6 @@ dependencies = [ "reqwest 0.12.28", "serde_json", "tokio", - "url", "wiremock", ] diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index 3890e60843c..4ec60bcfa3b 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -13,7 +13,6 @@ litellm-cache.workspace = true percent-encoding.workspace = true reqwest.workspace = true tokio.workspace = true -url.workspace = true [dev-dependencies] serde_json.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index 1097f3252b1..65282ac99d5 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -5,43 +5,18 @@ use litellm_cache::{ BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache, }; -use percent_encoding::{AsciiSet, CONTROLS, percent_encode}; +use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode}; use reqwest::Client; use crate::{GcpTokenSource, TokenSource}; pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; -const OBJECT_NAME_ENCODE_SET: &AsciiSet = &CONTROLS - .add(b' ') - .add(b'!') - .add(b'"') - .add(b'#') - .add(b'$') - .add(b'%') - .add(b'&') - .add(b'\'') - .add(b'(') - .add(b')') - .add(b'*') - .add(b'+') - .add(b',') - .add(b'/') - .add(b':') - .add(b';') - .add(b'<') - .add(b'=') - .add(b'>') - .add(b'?') - .add(b'@') - .add(b'[') - .add(b'\\') - .add(b']') - .add(b'^') - .add(b'`') - .add(b'{') - .add(b'|') - .add(b'}'); +const OBJECT_NAME_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC + .remove(b'-') + .remove(b'_') + .remove(b'.') + .remove(b'~'); pub fn key_prefix(gcs_path: Option<&str>) -> String { match gcs_path { diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs index 79e4ed100f8..45eecf01cec 100644 --- a/litellm-rust/crates/cache-gcs/tests/cache.rs +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -117,6 +117,40 @@ fn key_prefix_normalizes_paths() { assert_eq!(key_prefix(Some("")), ""); } +#[tokio::test] +async fn object_names_use_python_quote_encoding() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .respond_with(ResponseTemplate::new(200)) + .expect(2) + .mount(&server) + .await; + let cache = cache(&server, Some("p/")); + cache + .set_cache( + "a~b-c_d.e/f g%h", + json!({"value": "punctuation"}), + &ExactCacheContext::default(), + ) + .unwrap(); + cache + .set_cache( + "ключ", + json!({"value": "utf8"}), + &ExactCacheContext::default(), + ) + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let queries: Vec<_> = requests + .iter() + .filter_map(|request| request.url.query()) + .collect(); + assert!(queries.contains(&"uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h")); + assert!(queries.contains(&"uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87")); +} + #[tokio::test] async fn ignores_ttl_and_writes_pipeline_concurrently() { let server = MockServer::start().await; diff --git a/tests/test_litellm_rust/support/fake_gcs.py b/tests/test_litellm_rust/support/fake_gcs.py new file mode 100644 index 00000000000..67eb61798b9 --- /dev/null +++ b/tests/test_litellm_rust/support/fake_gcs.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import json +import threading +from collections.abc import Mapping +from dataclasses import dataclass +from functools import partial +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from socket import socket +from types import MappingProxyType +from typing import Final, cast +from urllib.parse import unquote, urlsplit + + +@dataclass(frozen=True, slots=True) +class RecordedRequest: + method: str + path: str + query: str + headers: Mapping[str, str] + body: bytes + + +class _FakeGcsHandler(BaseHTTPRequestHandler): + def __init__( + self, + request: socket | tuple[bytes, socket], + client_address: tuple[str, int], + server: ThreadingHTTPServer, + *, + fake: FakeGcs, + ) -> None: + self._fake: Final = fake + super().__init__(request, client_address, server) + + def _handle(self) -> None: + parsed: Final = urlsplit(self.path) + content_length: Final = int(self.headers.get("Content-Length", "0")) + body: Final = self.rfile.read(content_length) if content_length else b"" + headers: Final = MappingProxyType( + {name.title(): value for name, value in self.headers.items()} + ) + self._fake.record( + RecordedRequest( + method=self.command, + path=parsed.path, + query=parsed.query, + headers=headers, + body=body, + ) + ) + if self.headers.get("Authorization") != f"Bearer {self._fake.token}": + self._send_json(401, {"error": "unauthorized"}) + return + + upload_prefix: Final = "/upload/storage/v1/b/" + download_prefix: Final = "/storage/v1/b/" + if parsed.path.startswith(upload_prefix) and parsed.path.endswith("/o"): + self._upload(parsed.path[len(upload_prefix) : -2], parsed.query, body) + return + if parsed.path.startswith(download_prefix): + self._download(parsed.path[len(download_prefix) :], parsed.query) + return + self._send_json(404, {"error": "not found"}) + + def _upload(self, path: str, query: str, body: bytes) -> None: + values: Final = { + unquote(pair.partition("=")[0]): unquote(pair.partition("=")[2]) + for pair in query.split("&") + if pair + } + if not path or values.get("uploadType") != "media" or "name" not in values: + self._send_json(404, {"error": "not found"}) + return + self._fake.put_object(path, values["name"], body) + self._send_json(200, {"name": values["name"], "bucket": path}) + + def _download(self, path: str, query: str) -> None: + bucket, separator, encoded_name = path.partition("/o/") + if not separator or query != "alt=media": + self._send_json(404, {"error": "not found"}) + return + name: Final = unquote(encoded_name) + if name.endswith("/server-error") or name == "server-error": + self._send_json(500, {"error": "server error"}) + return + body: Final = self._fake.get_object(bucket, name) + if body is None: + self._send_json(404, {"error": "not found"}) + return + self._send(200, body, "application/octet-stream") + + def _send_json(self, status: int, value: object) -> None: + payload: Final = json.dumps(value).encode() + self._send(status, payload, "application/json") + + def _send(self, status: int, body: bytes, content_type: str) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + pass + + do_GET = _handle + do_POST = _handle + + +class FakeGcs: + def __init__(self) -> None: + self._objects: dict[tuple[str, str], bytes] = {} # mutable-ok: fake object store + self._requests: list[RecordedRequest] = [] # mutable-ok: recorded request history + self._server = ThreadingHTTPServer( + ("127.0.0.1", 0), + partial(_FakeGcsHandler, fake=self), + ) + self._worker = threading.Thread(target=self._server.serve_forever, daemon=True) + self._worker.start() + self.token: Final = "test-token" + + @property + def url(self) -> str: + address: Final = cast(tuple[str, int], self._server.server_address) + host, port = address + return f"http://{host}:{port}" + + @property + def objects(self) -> Mapping[tuple[str, str], bytes]: + return MappingProxyType(self._objects) + + @property + def requests(self) -> tuple[RecordedRequest, ...]: + return tuple(self._requests) + + def put(self, bucket: str, name: str, body: bytes) -> None: + self.put_object(bucket, name, body) + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + self._worker.join(timeout=5) + + def record(self, request: RecordedRequest) -> None: + self._requests.append(request) + + def put_object(self, bucket: str, name: str, body: bytes) -> None: + self._objects[(bucket, name)] = body + + def get_object(self, bucket: str, name: str) -> bytes | None: + return self._objects.get((bucket, name)) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index c35cb1a20fb..cc409ebe49c 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -16,9 +16,11 @@ import redis import litellm from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache +from litellm.caching.gcs_cache import GCSCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.fake_gcs import FakeGcs from tests.test_litellm_rust.support.isolation import rebound pytestmark: Final = pytest.mark.requires_rust_extension @@ -26,6 +28,7 @@ pytestmark: Final = pytest.mark.requires_rust_extension class CacheLookup(Protocol): def get_cache(self, **kwargs: object) -> object: ... + def flush_cache(self) -> object: ... def request(key: str = "key") -> dict[str, object]: @@ -45,6 +48,15 @@ def redis_url() -> Generator[str]: worker.join(timeout=5) +@pytest.fixture +def fake_gcs() -> Generator[FakeGcs]: + server: Final = FakeGcs() + try: + yield server + finally: + server.close() + + def test_existing_constructor_and_global_are_unchanged() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) assert type(facade.cache) is InMemoryCache @@ -393,3 +405,219 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: assert client.get("second") is not None await facade.cache.disconnect() client.close() + + +async def test_gcs_reads_python_entries_and_writes_python_compatible_objects( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + fake_gcs.put( + "bucket", + "cache/sync", + json.dumps({"timestamp": time.time(), "response": json.dumps(response)}).encode(), + ) + fake_gcs.put("bucket", "cache/async", json.dumps({"timestamp": time.time(), "response": response}).encode()) + fake_gcs.put("bucket", "cache/raw", json.dumps(response).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + assert binding.lookup(request("missing")) is None + + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = fake_gcs.objects[("bucket", "cache/native")] + stored_value: Final = cast(dict[str, object], json.loads(stored)) + assert stored_value["response"] == response + assert isinstance(stored_value["timestamp"], float) + upload: Final = next(item for item in fake_gcs.requests if item.method == "POST") + assert upload.path == "/upload/storage/v1/b/bucket/o" + assert upload.query == "uploadType=media&name=cache%2Fnative" + assert upload.headers["Authorization"] == f"Bearer {fake_gcs.token}" + assert upload.headers["Content-Type"] == "application/json" + upload_text: Final = f"{upload.path}?{upload.query}{upload.headers}" + assert "ttl" not in upload_text.lower() + assert "expiry" not in upload_text.lower() + download: Final = next(item for item in fake_gcs.requests if item.path.endswith("/cache%2Fsync")) + assert download.path == "/storage/v1/b/bucket/o/cache%2Fsync" + assert download.query == "alt=media" + + binding.store(request("sync2"), response) + assert binding.lookup(request("sync2")) == response + assert GCSCache(bucket_name="bucket", gcs_path="cache").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket", gcs_path="cache/").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket").key_prefix == "" + + +async def test_gcs_batch_lookup_preserves_order_and_treats_malformed_entries_as_misses(fake_gcs: FakeGcs) -> None: + fake_gcs.put("bucket", "cache/hit", json.dumps({"timestamp": time.time(), "response": {"value": 1}}).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + requests: Final = [request("hit"), request("missing"), request("invalid")] + expected: Final = {"values": [{"value": 1}, None, None], "missing_indices": [1, 2]} + + assert await binding.async_lookup_batch(requests) == expected + assert binding.lookup_batch(requests) == expected + await binding.async_store_batch([request("first"), request("second")], [{"value": 1}, {"value": 2}]) + assert ("bucket", "cache/first") in fake_gcs.objects + assert ("bucket", "cache/second") in fake_gcs.objects + + +async def test_gcs_facade_binds_only_exact_matching_configuration( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/nonexistent") + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + assert type(facade.cache) is GCSCache + + mismatched_bucket: Final = _native._CacheTestHandle.gcs( + "other", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="buckets must match"): + mismatched_bucket._bind_facade(facade) + mismatched_prefix: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="x", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="key prefixes must match"): + mismatched_prefix._bind_facade(facade) + mismatched_credentials: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + path_service_account="sa.json", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="credentials must match"): + mismatched_credentials._bind_facade(facade) + with pytest.raises(TypeError, match="types must match"): + _native._CacheTestHandle.memory()._bind_facade(facade) + + matching: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + matching._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + await binding.async_store(request("native"), {"value": "native"}) + assert await binding.async_lookup(request("native")) == {"value": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="native") is None + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "key_prefix", "x/"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "path_service_account", "sa.json"): + assert resolver.resolve().kind == "python_callback" + def no_get_cache(*args: object, **kwargs: object) -> None: + return None + + with rebound(facade.cache, "get_cache", no_get_cache): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + + class CustomGcs(GCSCache): + pass + + with rebound(facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + assert resolver.resolve().kind == "python_callback" + custom_facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with rebound(custom_facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + with pytest.raises(TypeError, match="types must match"): + matching._bind_facade(custom_facade) + + missing_bucket: Final = Cache(type=LiteLLMCacheType.GCS) + with pytest.raises(TypeError, match="requires a configured bucket name"): + matching._bind_facade(missing_bucket) + + +async def test_gcs_flush_is_a_no_op_and_ping_is_not_implemented( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + await binding.async_store(request("key"), {"value": "stored"}) + await binding.async_flush() + assert ("bucket", "cache/key") in fake_gcs.objects + assert await binding.async_lookup(request("key")) == {"value": "stored"} + with pytest.raises(NotImplementedError): + await binding.ping() + + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with pytest.raises(AttributeError): + await facade.ping() + assert cast(CacheLookup, facade.cache).flush_cache() is None + + +async def test_gcs_unauthorized_and_server_errors_surface_as_runtime_errors(fake_gcs: FakeGcs) -> None: + wrong_token: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token="wrong-token", + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + wrong_token.lookup(request("missing")) + assert not fake_gcs.objects + + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + binding.lookup(request("server-error")) + assert binding.lookup(request("missing")) is None From c7458bf5fe20bcbb7902fd259fd3b113b2b9e168 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:42:49 +0000 Subject: [PATCH 07/14] refactor(rust): build vault requests without local reassignment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../secrets-hashicorp/src/secret_manager.rs | 59 +++++++++++-------- 1 file changed, 34 insertions(+), 25 deletions(-) diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index e4ca8c387c3..52cc143c1e3 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -9,11 +9,11 @@ use litellm_secrets_types::{ BaseSecretManager, SecretValue, async_rotate_secret, validate_secret_name, }; use moka::future::Cache; -use reqwest::Client; +use reqwest::{Client, Identity, RequestBuilder}; use serde_json::{Value, json}; use tokio::sync::Mutex; -use crate::{Error, HashicorpVaultConfig}; +use crate::{Error, HashicorpVaultConfig, TlsCertAuth}; const CACHE_CAPACITY: u64 = 200; @@ -242,10 +242,10 @@ impl HashicorpVault { return Ok(token); } }; - let mut request = self.client.post(login_url).json(&body); - if let Some(namespace) = self.config.login_namespace() { - request = request.header("X-Vault-Namespace", namespace); - } + let request: RequestBuilder = with_namespace( + self.client.post(login_url).json(&body), + self.config.login_namespace(), + ); let response: reqwest::Response = request.send().await?; if !response.status().is_success() { return Err(Error::LoginStatus { @@ -305,24 +305,33 @@ impl BaseSecretManager for HashicorpVault { } fn client_for_config(config: &HashicorpVaultConfig) -> Result { - let mut builder: reqwest::ClientBuilder = Client::builder(); - if let Some(tls) = config.tls_cert.as_ref() { - let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { - path: tls.cert_path.clone(), - message: source.to_string(), - })?; - let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { - path: tls.key_path.clone(), - message: source.to_string(), - })?; - let identity: reqwest::Identity = reqwest::Identity::from_pem( - &[cert.as_slice(), key.as_slice()].concat(), - ) - .map_err(|source| Error::TlsIdentity { - path: tls.cert_path.clone(), - message: source.to_string(), - })?; - builder = builder.identity(identity); - } + let builder: reqwest::ClientBuilder = match config.tls_cert.as_ref() { + Some(tls) => Client::builder().identity(identity_for(tls)?), + None => Client::builder(), + }; builder.build().map_err(Error::Request) } + +fn with_namespace(request: RequestBuilder, namespace: Option<&str>) -> RequestBuilder { + match namespace { + Some(namespace) => request.header("X-Vault-Namespace", namespace), + None => request, + } +} + +fn identity_for(tls: &TlsCertAuth) -> Result { + let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + })?; + let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { + path: tls.key_path.clone(), + message: source.to_string(), + })?; + Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| { + Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + } + }) +} From 7c8f7d3f2c7d8356ac3662e153c40fd0a64e82fd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:18:16 +0000 Subject: [PATCH 08/14] fix(rust): reduce native wheel size Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 80425d567c0..3e9d2d2fbb5 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -76,7 +76,7 @@ veil = "0.3.0" [profile.release] opt-level = 3 -lto = "thin" +lto = "fat" codegen-units = 1 panic = "unwind" debug = false From ef2ab74c7aecab4f970827ec21d740ed3aaa430b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:33:38 +0000 Subject: [PATCH 09/14] refactor(rust): back the vault secret manager with vaultrs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 91 +++- litellm-rust/Cargo.toml | 3 + .../crates/secrets-hashicorp/Cargo.toml | 4 +- .../secrets-hashicorp/src/cert_login.rs | 25 ++ .../crates/secrets-hashicorp/src/error.rs | 8 +- .../crates/secrets-hashicorp/src/lib.rs | 3 +- .../secrets-hashicorp/src/secret_manager.rs | 406 +++++++++--------- .../secrets-hashicorp/tests/secret_manager.rs | 293 +++++++++---- litellm-rust/crates/secrets/tests/handler.rs | 39 +- 9 files changed, 586 insertions(+), 286 deletions(-) create mode 100644 litellm-rust/crates/secrets-hashicorp/src/cert_login.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 6f05bba9417..7b59533ad06 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2763,13 +2763,15 @@ dependencies = [ "litellm-core-utils", "litellm-secrets-types", "moka", - "reqwest 0.12.28", "rstest", + "rustify", + "rustify_derive", "serde", "serde_json", "tempfile", "thiserror 2.0.19", "tokio", + "vaultrs", "veil", "wiremock", ] @@ -3975,6 +3977,40 @@ dependencies = [ "semver", ] +[[package]] +name = "rustify" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4800ce4c1cc2fec12c559dae2ddbf0e17fcee7569b796e6d75898efef443368b" +dependencies = [ + "anyhow", + "async-trait", + "bytes", + "http 1.4.2", + "reqwest 0.13.5", + "rustify_derive", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror 1.0.69", + "tracing", + "url", +] + +[[package]] +name = "rustify_derive" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78ea7fda74240f7410d0198b603a8a2f662acc7d76b6667a49f9b162cd8d9b4f" +dependencies = [ + "proc-macro2", + "quote", + "regex", + "serde_urlencoded", + "syn 1.0.109", + "synstructure 0.12.6", +] + [[package]] name = "rustix" version = "1.1.5" @@ -4502,6 +4538,17 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "2.0.119" @@ -4533,6 +4580,18 @@ dependencies = [ "futures-core", ] +[[package]] +name = "synstructure" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f36bdaa60a83aca3921b5259d5400cbf5e90fc51931376a9bd4a0eb79aa7210f" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", + "unicode-xid", +] + [[package]] name = "synstructure" version = "0.13.2" @@ -4947,6 +5006,7 @@ version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ + "log", "pin-project-lite", "tracing-attributes", "tracing-core", @@ -5130,6 +5190,12 @@ version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + [[package]] name = "unicode_categories" version = "0.1.1" @@ -5189,6 +5255,25 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "vaultrs" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30ffcc0e81025065dda612ec1e26a3d81bb16ef3062354873d17a35965d68522" +dependencies = [ + "async-trait", + "derive_builder", + "http 1.4.2", + "reqwest 0.13.5", + "rustify", + "rustify_derive", + "serde", + "serde_json", + "thiserror 2.0.19", + "tracing", + "url", +] + [[package]] name = "veil" version = "0.3.0" @@ -5632,7 +5717,7 @@ dependencies = [ "proc-macro2", "quote", "syn 2.0.119", - "synstructure", + "synstructure 0.13.2", ] [[package]] @@ -5673,7 +5758,7 @@ dependencies = [ "proc-macro2", "quote", "syn 2.0.119", - "synstructure", + "synstructure 0.13.2", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 5030a94b140..edc9f5a91f6 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -51,6 +51,9 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "mul rstest = "0.26.1" rstest_reuse = "0.7.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } +rustify = "=0.7.0" +rustify_derive = "=0.5.5" +vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] } rustls-native-certs = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml index 0656980a033..c646e02ef09 100644 --- a/litellm-rust/crates/secrets-hashicorp/Cargo.toml +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -9,11 +9,13 @@ repository.workspace = true litellm-core-utils.workspace = true litellm-secrets-types.workspace = true moka.workspace = true -reqwest.workspace = true +rustify.workspace = true +rustify_derive.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true tokio.workspace = true +vaultrs.workspace = true veil.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs b/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs new file mode 100644 index 00000000000..1eb09a89715 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs @@ -0,0 +1,25 @@ +#[derive(Debug, rustify_derive::Endpoint)] +#[endpoint(path = "/auth/{self.mount}/login", method = "POST")] +pub struct CertLoginRequest { + #[endpoint(skip)] + pub mount: String, + #[endpoint(skip)] + #[allow(dead_code)] + pub name: Option, + #[endpoint(raw)] + body: Vec, +} + +impl CertLoginRequest { + pub fn new(name: Option) -> Self { + let body: Vec = match name.as_deref() { + Some(name) => serde_json::to_vec(&serde_json::json!({ "name": name })).unwrap(), + None => b"{}".to_vec(), + }; + Self { + mount: "cert".to_owned(), + name, + body, + } + } +} diff --git a/litellm-rust/crates/secrets-hashicorp/src/error.rs b/litellm-rust/crates/secrets-hashicorp/src/error.rs index 3f6a5b66475..e26033af085 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/error.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/error.rs @@ -4,12 +4,14 @@ pub enum Error { EnterpriseRequired, #[error("invalid secret name")] InvalidSecretName(#[from] litellm_secrets_types::Error), - #[error("HashiCorp Vault request failed")] - Request( + #[error("HashiCorp Vault client failed")] + Client( #[from] #[redact] - reqwest::Error, + vaultrs::error::ClientError, ), + #[error("HashiCorp Vault client settings are invalid: {message}")] + ClientSettings { message: String }, #[error("HashiCorp Vault TLS identity could not be configured for {path}: {message}")] TlsIdentity { path: std::path::PathBuf, diff --git a/litellm-rust/crates/secrets-hashicorp/src/lib.rs b/litellm-rust/crates/secrets-hashicorp/src/lib.rs index a561a58b59e..0c2b05647f8 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/lib.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/lib.rs @@ -1,9 +1,10 @@ #![forbid(unsafe_code)] +mod cert_login; mod config; mod error; pub mod secret_manager; pub use config::{AppRoleAuth, HashicorpVaultConfig, TlsCertAuth}; pub use error::Error; -pub use secret_manager::HashicorpVault; +pub use secret_manager::{HashicorpVault, SecretLocation}; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index 52cc143c1e3..7937fa93f8f 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -1,4 +1,5 @@ use std::{ + collections::HashMap, fmt, sync::Arc, time::{Duration, Instant}, @@ -9,26 +10,39 @@ use litellm_secrets_types::{ BaseSecretManager, SecretValue, async_rotate_secret, validate_secret_name, }; use moka::future::Cache; -use reqwest::{Client, Identity, RequestBuilder}; -use serde_json::{Value, json}; +use rustify::errors::ClientError as RustifyClientError; +use serde_json::Value; use tokio::sync::Mutex; +use vaultrs::{ + api, + auth::approle, + client::{Identity, VaultClient, VaultClientSettingsBuilder}, + error::ClientError, + kv2, +}; -use crate::{Error, HashicorpVaultConfig, TlsCertAuth}; +use crate::{Error, HashicorpVaultConfig, TlsCertAuth, cert_login::CertLoginRequest}; const CACHE_CAPACITY: u64 = 200; #[derive(Clone)] -struct CachedToken { - token: SecretValue, +struct CachedClient { + client: Arc, expires_at: Option, } +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SecretLocation { + pub namespace: Option, + pub mount: String, + pub path: String, +} + #[derive(Clone)] pub struct HashicorpVault { - client: Arc, config: HashicorpVaultConfig, cache: Cache, - auth_token: Arc>>, + auth_client: Arc>>, } impl fmt::Debug for HashicorpVault { @@ -45,17 +59,12 @@ impl HashicorpVault { environment: Arc, enterprise_enabled: bool, ) -> Result { - if !enterprise_enabled { - return Err(Error::EnterpriseRequired); - } let config: HashicorpVaultConfig = HashicorpVaultConfig::from_environment(environment.as_ref())?; - let client: Client = client_for_config(&config)?; - Self::with_client(client, config, enterprise_enabled) + Self::from_config(config, enterprise_enabled) } - pub fn with_client( - client: Client, + pub fn from_config( config: HashicorpVaultConfig, enterprise_enabled: bool, ) -> Result { @@ -67,47 +76,27 @@ impl HashicorpVault { .time_to_live(config.refresh_interval) .build(); Ok(Self { - client: Arc::new(client), config, cache, - auth_token: Arc::new(Mutex::new(None)), + auth_client: Arc::new(Mutex::new(None)), }) } - pub fn secret_url(&self, secret_name: &str) -> Result { + pub fn secret_location(&self, secret_name: &str) -> Result { validate_secret_name(secret_name).map_err(Error::InvalidSecretName)?; - let namespace: String = self - .config - .secret_namespace() - .map(|value| format!("{value}/")) - .unwrap_or_default(); - let path_prefix: String = self - .config - .path_prefix - .as_deref() - .map(|value| format!("{value}/")) - .unwrap_or_default(); - Ok(format!( - "{}/v1/{}{}/data/{}{}", - self.config.address, namespace, self.config.mount, path_prefix, secret_name - )) - } - - pub fn login_url(&self) -> Option { - self.config.approle.as_ref().map_or_else( - || { - self.config - .tls_cert - .as_ref() - .map(|_| format!("{}/v1/auth/cert/login", self.config.address)) - }, - |approle| { - Some(format!( - "{}/v1/auth/{}/login", - self.config.address, approle.mount_path - )) - }, - ) + let path: String = [ + self.config.path_prefix.clone(), + Some(secret_name.to_owned()), + ] + .into_iter() + .flatten() + .collect::>() + .join("/"); + Ok(SecretLocation { + namespace: self.config.secret_namespace().map(str::to_owned), + mount: self.config.mount.clone(), + path, + }) } pub fn config(&self) -> &HashicorpVaultConfig { @@ -115,40 +104,24 @@ impl HashicorpVault { } pub async fn async_read_secret(&self, secret_name: &str) -> Result, Error> { - let url: String = self.secret_url(secret_name)?; - if let Some(value) = self.cache.get(&url).await { + let location: SecretLocation = self.secret_location(secret_name)?; + let cache_key: String = cache_key(&location); + if let Some(value) = self.cache.get(&cache_key).await { return Ok(Some(value)); } - let token: SecretValue = self.vault_token().await?; - let response: reqwest::Response = self - .client - .get(&url) - .header("X-Vault-Token", token.expose()) - .send() - .await?; - if response.status() == reqwest::StatusCode::NOT_FOUND { - return Ok(None); - } - if !response.status().is_success() { - return Err(Error::Status { - status: response.status().as_u16(), - }); - } - let body = response.bytes().await?; - let body: Value = serde_json::from_slice(&body).map_err(|_| Error::MalformedPayload)?; - let data: &Value = body - .get("data") - .and_then(Value::as_object) - .and_then(|value| value.get("data")) - .ok_or(Error::MalformedPayload)?; - let data: &serde_json::Map = - data.as_object().ok_or(Error::MalformedPayload)?; + let client: Arc = self.vault_client().await?; + let data: HashMap = + match kv2::read(client.as_ref(), &location.mount, &location.path).await { + Ok(data) => data, + Err(error) if api_status(&error) == Some(404) => return Ok(None), + Err(error) => return Err(map_api_error(error, ErrorContext::Read)), + }; let Some(value) = data.get("key") else { return Ok(None); }; let value: &str = value.as_str().ok_or(Error::NonStringValue)?; let value: SecretValue = SecretValue::new(value); - self.cache.insert(url, value.clone()).await; + self.cache.insert(cache_key, value.clone()).await; Ok(Some(value)) } @@ -158,47 +131,39 @@ impl HashicorpVault { value: SecretValue, description: Option<&str>, ) -> Result { - let url: String = self.secret_url(secret_name)?; - let data: Value = match description { - Some(description) => json!({"key": value.expose(), "description": description}), - None => json!({"key": value.expose()}), + let location: SecretLocation = self.secret_location(secret_name)?; + let cache_key: String = cache_key(&location); + let data: HashMap = match description { + Some(description) => [ + ("key".to_owned(), Value::String(value.expose().to_owned())), + ( + "description".to_owned(), + Value::String(description.to_owned()), + ), + ] + .into_iter() + .collect(), + None => [("key".to_owned(), Value::String(value.expose().to_owned()))] + .into_iter() + .collect(), }; - let token: SecretValue = self.vault_token().await?; - let response: reqwest::Response = self - .client - .post(&url) - .header("X-Vault-Token", token.expose()) - .json(&json!({"data": data})) - .send() - .await?; - if !response.status().is_success() { - return Err(Error::Status { - status: response.status().as_u16(), - }); - } - self.cache.invalidate(&url).await; - let body = response.bytes().await?; - if body.is_empty() { - return Ok(Value::Null); - } - serde_json::from_slice(&body).map_err(|_| Error::MalformedPayload) + let client: Arc = self.vault_client().await?; + let metadata = kv2::set(client.as_ref(), &location.mount, &location.path, &data) + .await + .map_err(|error| map_api_error(error, ErrorContext::Secret))?; + self.cache.invalidate(&cache_key).await; + serde_json::to_value(metadata) + .map_err(|source| Error::Client(ClientError::JsonParseError { source })) } pub async fn async_delete_secret(&self, secret_name: &str) -> Result<(), Error> { - let url: String = self.secret_url(secret_name)?; - let token: SecretValue = self.vault_token().await?; - let response: reqwest::Response = self - .client - .delete(&url) - .header("X-Vault-Token", token.expose()) - .send() - .await?; - if !response.status().is_success() { - return Err(Error::Status { - status: response.status().as_u16(), - }); - } - self.cache.invalidate(&url).await; + let location: SecretLocation = self.secret_location(secret_name)?; + let cache_key: String = cache_key(&location); + let client: Arc = self.vault_client().await?; + kv2::delete_latest(client.as_ref(), &location.mount, &location.path) + .await + .map_err(|error| map_api_error(error, ErrorContext::Secret))?; + self.cache.invalidate(&cache_key).await; Ok(()) } @@ -211,69 +176,76 @@ impl HashicorpVault { async_rotate_secret(self, current_name, new_name, value).await } - async fn vault_token(&self) -> Result { - let mut cached: tokio::sync::MutexGuard<'_, Option> = - self.auth_token.lock().await; + async fn vault_client(&self) -> Result, Error> { + let mut cached = self.auth_client.lock().await; if let Some(entry) = cached.as_ref() && entry .expires_at .is_none_or(|expires_at| expires_at > Instant::now()) { - return Ok(entry.token.clone()); + return Ok(entry.client.clone()); } - let Some(login_url) = self.login_url() else { - let Some(token) = self.config.token.clone() else { - return Err(Error::NoAuthConfigured); + + let (client, expires_at): (VaultClient, Option) = + match (self.config.approle.as_ref(), self.config.tls_cert.as_ref()) { + (Some(approle), _) => { + let login_client: VaultClient = + self.build_client(self.config.login_namespace(), "")?; + let auth = approle::login( + &login_client, + &approle.mount_path, + &approle.role_id, + approle.secret_id.expose(), + ) + .await + .map_err(|error| map_api_error(error, ErrorContext::Login))?; + ( + self.build_client(self.config.secret_namespace(), &auth.client_token)?, + token_expiry(auth.lease_duration), + ) + } + (None, Some(tls)) => { + let login_client: VaultClient = + self.build_client(self.config.login_namespace(), "")?; + let endpoint: CertLoginRequest = CertLoginRequest::new(tls.role.clone()); + let auth = api::auth(&login_client, endpoint) + .await + .map_err(|error| map_api_error(error, ErrorContext::Login))?; + ( + self.build_client(self.config.secret_namespace(), &auth.client_token)?, + token_expiry(auth.lease_duration), + ) + } + (None, None) => { + let token: SecretValue = + self.config.token.clone().ok_or(Error::NoAuthConfigured)?; + ( + self.build_client(self.config.secret_namespace(), token.expose())?, + None, + ) + } }; - return Ok(token); - }; - let body: Value = match (self.config.approle.as_ref(), self.config.tls_cert.as_ref()) { - (Some(approle), _) => { - json!({"role_id": approle.role_id, "secret_id": approle.secret_id.expose()}) - } - (None, Some(tls)) => tls - .role - .as_deref() - .map_or_else(|| json!({}), |role| json!({"name": role})), - (None, None) => { - let Some(token) = self.config.token.clone() else { - return Err(Error::NoAuthConfigured); - }; - return Ok(token); - } - }; - let request: RequestBuilder = with_namespace( - self.client.post(login_url).json(&body), - self.config.login_namespace(), - ); - let response: reqwest::Response = request.send().await?; - if !response.status().is_success() { - return Err(Error::LoginStatus { - status: response.status().as_u16(), - }); - } - let body = response.bytes().await?; - let payload: Value = serde_json::from_slice(&body).map_err(|_| Error::MalformedLogin)?; - let auth: &serde_json::Map = payload - .get("auth") - .and_then(Value::as_object) - .ok_or(Error::MalformedLogin)?; - let token: SecretValue = SecretValue::new( - auth.get("client_token") - .and_then(Value::as_str) - .ok_or(Error::MalformedLogin)?, - ); - let lease_duration: u64 = auth - .get("lease_duration") - .and_then(Value::as_u64) - .ok_or(Error::MalformedLogin)?; - let expires_at: Option = - (lease_duration > 0).then(|| Instant::now() + Duration::from_secs(lease_duration)); - *cached = Some(CachedToken { - token: token.clone(), + let client: Arc = Arc::new(client); + *cached = Some(CachedClient { + client: client.clone(), expires_at, }); - Ok(token) + Ok(client) + } + + fn build_client(&self, namespace: Option<&str>, token: &str) -> Result { + let settings = VaultClientSettingsBuilder::default() + .address(&self.config.address) + .token(token.to_owned()) + .namespace(namespace.map(str::to_owned)) + .identity(identity_for(self.config.tls_cert.as_ref())?) + .ca_certs(Vec::new()) + .verify(true) + .build() + .map_err(|message| Error::ClientSettings { + message: message.to_string(), + })?; + VaultClient::new(settings).map_err(Error::Client) } } @@ -304,34 +276,84 @@ impl BaseSecretManager for HashicorpVault { } } -fn client_for_config(config: &HashicorpVaultConfig) -> Result { - let builder: reqwest::ClientBuilder = match config.tls_cert.as_ref() { - Some(tls) => Client::builder().identity(identity_for(tls)?), - None => Client::builder(), - }; - builder.build().map_err(Error::Request) +#[derive(Clone, Copy)] +enum ErrorContext { + Login, + Read, + Secret, } -fn with_namespace(request: RequestBuilder, namespace: Option<&str>) -> RequestBuilder { - match namespace { - Some(namespace) => request.header("X-Vault-Namespace", namespace), - None => request, +fn cache_key(location: &SecretLocation) -> String { + format!( + "{:?}/{}/{}", + location.namespace, location.mount, location.path + ) +} + +fn identity_for(tls: Option<&TlsCertAuth>) -> Result, Error> { + tls.map(|tls| { + let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + })?; + let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { + path: tls.key_path.clone(), + message: source.to_string(), + })?; + Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| { + Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + } + }) + }) + .transpose() +} + +fn map_api_error(error: ClientError, context: ErrorContext) -> Error { + match error { + ClientError::APIError { code, .. } => match context { + ErrorContext::Login => Error::LoginStatus { status: code }, + ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code }, + }, + ClientError::JsonParseError { source } => match context { + ErrorContext::Login => Error::MalformedLogin, + ErrorContext::Read => Error::MalformedPayload, + ErrorContext::Secret => Error::Client(ClientError::JsonParseError { source }), + }, + ClientError::ResponseEmptyError | ClientError::ResponseDataEmptyError => { + malformed_response(context) + } + ClientError::RestClientError { source } => match source { + RustifyClientError::ServerResponseError { code, .. } => match context { + ErrorContext::Login => Error::LoginStatus { status: code }, + ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code }, + }, + RustifyClientError::ResponseParseError { .. } => malformed_response(context), + source => Error::Client(ClientError::RestClientError { source }), + }, + error => Error::Client(error), } } -fn identity_for(tls: &TlsCertAuth) -> Result { - let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { - path: tls.cert_path.clone(), - message: source.to_string(), - })?; - let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { - path: tls.key_path.clone(), - message: source.to_string(), - })?; - Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| { - Error::TlsIdentity { - path: tls.cert_path.clone(), - message: source.to_string(), - } - }) +fn api_status(error: &ClientError) -> Option { + match error { + ClientError::APIError { code, .. } => Some(*code), + ClientError::RestClientError { + source: RustifyClientError::ServerResponseError { code, .. }, + } => Some(*code), + _ => None, + } +} + +fn malformed_response(context: ErrorContext) -> Error { + match context { + ErrorContext::Login => Error::MalformedLogin, + ErrorContext::Read => Error::MalformedPayload, + ErrorContext::Secret => Error::Client(ClientError::ResponseDataEmptyError), + } +} + +fn token_expiry(lease_duration: u64) -> Option { + (lease_duration > 0).then(|| Instant::now() + Duration::from_secs(lease_duration)) } diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs index b7a81f49ad4..c52db46e41e 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -22,7 +22,51 @@ fn config(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVaultConfig } fn manager(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVault { - HashicorpVault::with_client(reqwest::Client::new(), config(server, values), true).unwrap() + HashicorpVault::from_config(config(server, values), true).unwrap() +} + +fn auth_response(token: &str, lease_duration: u64) -> serde_json::Value { + json!({ + "auth": { + "client_token": token, + "accessor": "", + "policies": [], + "token_policies": [], + "metadata": null, + "lease_duration": lease_duration, + "renewable": false, + "entity_id": "", + "token_type": "service", + "orphan": false + }, + "lease_id": "", + "lease_duration": lease_duration, + "renewable": false, + "request_id": "", + "warnings": null, + "wrap_info": null + }) +} + +fn read_response(data: serde_json::Value) -> serde_json::Value { + json!({ + "data": { + "data": data, + "metadata": { + "created_time": "", + "deletion_time": "", + "custom_metadata": null, + "destroyed": false, + "version": 1 + } + }, + "lease_id": "", + "lease_duration": 0, + "renewable": false, + "request_id": "", + "warnings": null, + "wrap_info": null + }) } #[tokio::test] @@ -32,7 +76,7 @@ async fn token_reads_use_vault_headers_and_cache_values() { .and(path("/v1/secret/data/name")) .and(header("X-Vault-Token", "token")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(1) .mount(&server) @@ -48,6 +92,12 @@ async fn token_reads_use_vault_headers_and_cache_values() { .expose(), "value" ); + let requests = server.received_requests().await.unwrap(); + assert!( + requests + .iter() + .all(|request| !request.headers.contains_key("X-Vault-Namespace")) + ); assert_eq!( manager .async_read_secret("name") @@ -63,9 +113,10 @@ async fn token_reads_use_vault_headers_and_cache_values() { async fn namespace_mount_and_prefix_are_sanitized_in_the_url() { let server: MockServer = MockServer::start().await; Mock::given(method("GET")) - .and(path("/v1/team-a/kv-prod/data/virtual-keys/name")) + .and(path("/v1/kv-prod/data/virtual-keys/name")) + .and(header("X-Vault-Namespace", "team-a")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(1) .mount(&server) @@ -80,10 +131,10 @@ async fn namespace_mount_and_prefix_are_sanitized_in_the_url() { ], ); - assert_eq!( - manager.secret_url("name").unwrap(), - format!("{}/v1/team-a/kv-prod/data/virtual-keys/name", server.uri()) - ); + let location = manager.secret_location("name").unwrap(); + assert_eq!(location.namespace.as_deref(), Some("team-a")); + assert_eq!(location.mount, "kv-prod"); + assert_eq!(location.path, "virtual-keys/name"); assert!(manager.async_read_secret("name").await.unwrap().is_some()); } @@ -96,12 +147,15 @@ fn trailing_address_slashes_are_removed() { }); let config: HashicorpVaultConfig = HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap(); - let manager: HashicorpVault = - HashicorpVault::with_client(reqwest::Client::new(), config, true).unwrap(); + let manager: HashicorpVault = HashicorpVault::from_config(config, true).unwrap(); assert_eq!( - manager.secret_url("name").unwrap(), - "http://vault.test:8200/v1/secret/data/name" + manager.secret_location("name").unwrap(), + litellm_secrets_hashicorp::SecretLocation { + namespace: None, + mount: "secret".to_owned(), + path: "name".to_owned(), + } ); } @@ -127,21 +181,26 @@ async fn approle_login_uses_namespace_and_reuses_the_token() { .and(path("/v1/auth/custom-approle/login")) .and(header("X-Vault-Namespace", "login-root")) .and(body_json(json!({"role_id": "role", "secret_id": "secret"}))) - .respond_with(ResponseTemplate::new(200).set_body_json( - json!({"auth": {"client_token": "login-token", "lease_duration": 3600}}), - )) + .respond_with(ResponseTemplate::new(200).set_body_json(auth_response("login-token", 3600))) .expect(1) .mount(&server) .await; Mock::given(method("GET")) - .and(path("/v1/secret-root/secret/data/name")) + .and(path("/v1/secret/data/name")) .and(header("X-Vault-Token", "login-token")) + .and(header("X-Vault-Namespace", "secret-root")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(1) .mount(&server) .await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/name-2")) + .respond_with(ResponseTemplate::new(404).set_body_json(json!({"errors": ["missing"]}))) + .expect(1) + .mount(&server) + .await; let manager: HashicorpVault = manager( &server, &[ @@ -162,17 +221,13 @@ async fn approle_tokens_expire_after_the_vault_lease() { let server: MockServer = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/auth/approle/login")) - .respond_with( - ResponseTemplate::new(200).set_body_json( - json!({"auth": {"client_token": "login-token", "lease_duration": 1}}), - ), - ) + .respond_with(ResponseTemplate::new(200).set_body_json(auth_response("login-token", 1))) .expect(2) .mount(&server) .await; Mock::given(method("GET")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(2) .mount(&server) @@ -201,54 +256,98 @@ async fn tls_login_posts_the_role_and_uses_the_client_identity() { std::fs::write(&key_path, TEST_PRIVATE_KEY).unwrap(); Mock::given(method("POST")) .and(path("/v1/auth/cert/login")) - .and(body_json(json!({"name": "vault-role"}))) - .respond_with( - ResponseTemplate::new(200).set_body_json( - json!({"auth": {"client_token": "cert-token", "lease_duration": 0}}), - ), - ) - .expect(1) + .and(header("X-Vault-Namespace", "login-ns")) + .respond_with(ResponseTemplate::new(200).set_body_json(auth_response("cert-token", 0))) + .expect(2) .mount(&server) .await; Mock::given(method("GET")) .and(path("/v1/secret/data/name")) + .and(header("X-Vault-Token", "cert-token")) + .and(header("X-Vault-Namespace", "secret-ns")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) + .expect(2) .mount(&server) .await; - let manager: HashicorpVault = HashicorpVault::new( - { - let environment_values: HashMap = HashMap::from([ - ("HCP_VAULT_ADDR".to_owned(), server.uri()), - ( - "HCP_VAULT_CLIENT_CERT".to_owned(), - cert_path.to_str().unwrap().to_owned(), - ), - ( - "HCP_VAULT_CLIENT_KEY".to_owned(), - key_path.to_str().unwrap().to_owned(), - ), - ("HCP_VAULT_CERT_ROLE".to_owned(), "vault-role".to_owned()), - ]); - Arc::new(move |name: &str| environment_values.get(name).cloned()) - }, - true, - ) - .unwrap(); - - assert!(manager.async_read_secret("name").await.unwrap().is_some()); - assert_eq!( - manager.login_url().as_deref(), - Some(format!("{}/v1/auth/cert/login", server.uri()).as_str()) + let role_values: HashMap = HashMap::from([ + ("HCP_VAULT_ADDR".to_owned(), server.uri()), + ( + "HCP_VAULT_CLIENT_CERT".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_CLIENT_KEY".to_owned(), + key_path.to_str().unwrap().to_owned(), + ), + ("HCP_VAULT_CERT_ROLE".to_owned(), "vault-role".to_owned()), + ( + "HCP_VAULT_LOGIN_NAMESPACE".to_owned(), + "login-ns".to_owned(), + ), + ( + "HCP_VAULT_SECRET_NAMESPACE".to_owned(), + "secret-ns".to_owned(), + ), + ]); + let role_environment: Arc = + Arc::new(move |name: &str| role_values.get(name).cloned()); + let role_manager: HashicorpVault = HashicorpVault::new(role_environment, true).unwrap(); + assert!( + role_manager + .async_read_secret("name") + .await + .unwrap() + .is_some() ); + + let no_role_values: HashMap = HashMap::from([ + ("HCP_VAULT_ADDR".to_owned(), server.uri()), + ( + "HCP_VAULT_CLIENT_CERT".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_CLIENT_KEY".to_owned(), + key_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_LOGIN_NAMESPACE".to_owned(), + "login-ns".to_owned(), + ), + ( + "HCP_VAULT_SECRET_NAMESPACE".to_owned(), + "secret-ns".to_owned(), + ), + ]); + let no_role_environment: Arc = + Arc::new(move |name: &str| no_role_values.get(name).cloned()); + let no_role_manager: HashicorpVault = HashicorpVault::new(no_role_environment, true).unwrap(); + assert!( + no_role_manager + .async_read_secret("name") + .await + .unwrap() + .is_some() + ); + let login_bodies: Vec = server + .received_requests() + .await + .unwrap() + .iter() + .filter(|request| request.method.as_str() == "POST") + .map(|request| serde_json::from_slice(&request.body).unwrap()) + .collect(); + assert!(login_bodies.contains(&json!({"name": "vault-role"}))); + assert!(login_bodies.contains(&json!({}))); } #[rstest::rstest] -#[case::missing(404, json!({}), 0)] -#[case::malformed(200, json!({}), 1)] -#[case::missing_key(200, json!({"data": {"data": {}}}), 0)] -#[case::non_string(200, json!({"data": {"data": {"key": 1}}}), 2)] +#[case::missing(404, json!({"errors": ["missing"]}), 0)] +#[case::malformed(200, json!({"data": "invalid"}), 1)] +#[case::missing_key(200, json!({}), 0)] +#[case::non_string(200, json!({"key": 1}), 2)] #[tokio::test] async fn read_responses_distinguish_absence_and_malformed_payloads( #[case] status: u16, @@ -257,7 +356,13 @@ async fn read_responses_distinguish_absence_and_malformed_payloads( ) { let server: MockServer = MockServer::start().await; Mock::given(method("GET")) - .respond_with(ResponseTemplate::new(status).set_body_json(body)) + .respond_with(ResponseTemplate::new(status).set_body_json( + if status == 200 && expected != 1 { + read_response(body) + } else { + body + }, + )) .expect(1) .mount(&server) .await; @@ -279,7 +384,7 @@ async fn write_and_delete_invalidate_the_read_cache() { Mock::given(method("GET")) .and(path("/v1/secret/data/name")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(2) .mount(&server) @@ -289,7 +394,21 @@ async fn write_and_delete_invalidate_the_read_cache() { .and(body_json( json!({"data": {"key": "updated", "description": "description"}}), )) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data": {"version": 2}}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "data": { + "created_time": "", + "deletion_time": "", + "custom_metadata": null, + "destroyed": false, + "version": 2 + }, + "lease_id": "", + "lease_duration": 0, + "renewable": false, + "request_id": "", + "warnings": null, + "wrap_info": null + }))) .expect(1) .mount(&server) .await; @@ -331,12 +450,9 @@ async fn no_auth_and_invalid_names_fail_without_requests() { #[tokio::test] async fn debug_output_redacts_authentication_values() { let server: MockServer = MockServer::start().await; - let manager: HashicorpVault = HashicorpVault::with_client( - reqwest::Client::new(), - config(&server, &[("HCP_VAULT_TOKEN", "token-value")]), - true, - ) - .unwrap(); + let manager: HashicorpVault = + HashicorpVault::from_config(config(&server, &[("HCP_VAULT_TOKEN", "token-value")]), true) + .unwrap(); let debug: String = format!("{manager:?}"); assert!(!debug.contains("token-value")); assert!(!debug.contains("secret-id")); @@ -365,13 +481,35 @@ fn configuration_matches_python_parity_fixture() { Arc::new(move |name: &str| values.get(name).cloned()); let config: HashicorpVaultConfig = HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap(); - let manager: HashicorpVault = - HashicorpVault::with_client(reqwest::Client::new(), config, true).unwrap(); + let manager: HashicorpVault = HashicorpVault::from_config(config.clone(), true).unwrap(); + let location = manager.secret_location(&case.secret_name).unwrap(); + let namespace = location + .namespace + .as_deref() + .map(|namespace| format!("{namespace}/")) + .unwrap_or_default(); assert_eq!( - manager.secret_url(&case.secret_name).unwrap(), + format!( + "{}/v1/{}{}/data/{}", + config.address, namespace, location.mount, location.path + ), case.expected_secret_url ); - assert_eq!(manager.login_url(), case.expected_login_url); + let login_url = config.approle.as_ref().map_or_else( + || { + config + .tls_cert + .as_ref() + .map(|_| format!("{}/v1/auth/cert/login", config.address)) + }, + |approle| { + Some(format!( + "{}/v1/auth/{}/login", + config.address, approle.mount_path + )) + }, + ); + assert_eq!(login_url, case.expected_login_url); assert_eq!( manager.config().login_namespace(), case.expected_login_namespace.as_deref() @@ -391,8 +529,15 @@ async fn live_vault_round_trip() { let manager: HashicorpVault = HashicorpVault::new(environment, true).unwrap(); let name: String = std::env::var("LITELLM_VAULT_LIVE_SECRET_NAME").unwrap(); let value: SecretValue = SecretValue::new("native-live-value"); - let url: String = manager.secret_url(&name).unwrap(); - println!("native provenance: {} {}", module_path!(), url); + let location = manager.secret_location(&name).unwrap(); + println!( + "native provenance: {} vaultrs {} {:?} {} {}", + module_path!(), + manager.config().address, + location.namespace, + location.mount, + location.path + ); manager .async_write_secret(&name, value.clone(), None) .await diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs index 58941171fd9..c169d411f8f 100644 --- a/litellm-rust/crates/secrets/tests/handler.rs +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -124,10 +124,24 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { let found_server = MockServer::start().await; Mock::given(method("GET")) .and(path("/v1/secret/data/KEY")) - .respond_with( - ResponseTemplate::new(200) - .set_body_json(serde_json::json!({"data": {"data": {"key": "remote"}}})), - ) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "data": { + "data": {"key": "remote"}, + "metadata": { + "created_time": "", + "deletion_time": "", + "custom_metadata": null, + "destroyed": false, + "version": 1 + } + }, + "lease_id": "", + "lease_duration": 0, + "renewable": false, + "request_id": "", + "warnings": null, + "wrap_info": null + }))) .mount(&found_server) .await; let found_environment: Arc = Arc::new({ @@ -139,8 +153,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { } }); let found_config = HashicorpVaultConfig::from_environment(found_environment.as_ref()).unwrap(); - let found_manager = - HashicorpVault::with_client(reqwest::Client::new(), found_config, true).unwrap(); + let found_manager = HashicorpVault::from_config(found_config, true).unwrap(); let found_resolver = SecretResolver::new( Arc::new(SecretManagerState::new( SecretManager::HashicorpVault(found_manager), @@ -164,7 +177,9 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { let missing_server = MockServer::start().await; Mock::given(method("GET")) - .respond_with(ResponseTemplate::new(404)) + .respond_with( + ResponseTemplate::new(404).set_body_json(serde_json::json!({"errors": ["missing"]})), + ) .mount(&missing_server) .await; let missing_environment: Arc = Arc::new({ @@ -177,8 +192,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { }); let missing_config = HashicorpVaultConfig::from_environment(missing_environment.as_ref()).unwrap(); - let missing_manager = - HashicorpVault::with_client(reqwest::Client::new(), missing_config, true).unwrap(); + let missing_manager = HashicorpVault::from_config(missing_config, true).unwrap(); let missing_state = SecretManagerState::new( SecretManager::HashicorpVault(missing_manager), KeyManagementSettings { @@ -198,7 +212,9 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { let failed_server = MockServer::start().await; Mock::given(method("GET")) - .respond_with(ResponseTemplate::new(500)) + .respond_with( + ResponseTemplate::new(500).set_body_json(serde_json::json!({"errors": ["failed"]})), + ) .mount(&failed_server) .await; let failed_environment: Arc = Arc::new({ @@ -211,8 +227,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { }); let failed_config = HashicorpVaultConfig::from_environment(failed_environment.as_ref()).unwrap(); - let failed_manager = - HashicorpVault::with_client(reqwest::Client::new(), failed_config, true).unwrap(); + let failed_manager = HashicorpVault::from_config(failed_config, true).unwrap(); let failed_state = SecretManagerState::new( SecretManager::HashicorpVault(failed_manager), KeyManagementSettings { From d90ff4a30b67b41199e171ee997d93cdea4b3a57 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:35:39 +0000 Subject: [PATCH 10/14] refactor(rust): trim vault cert login endpoint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/secrets-hashicorp/src/cert_login.rs | 11 ++++------- .../crates/secrets-hashicorp/src/secret_manager.rs | 2 +- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs b/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs index 1eb09a89715..f99df62db12 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs @@ -3,22 +3,19 @@ pub struct CertLoginRequest { #[endpoint(skip)] pub mount: String, - #[endpoint(skip)] - #[allow(dead_code)] - pub name: Option, #[endpoint(raw)] body: Vec, } impl CertLoginRequest { - pub fn new(name: Option) -> Self { - let body: Vec = match name.as_deref() { - Some(name) => serde_json::to_vec(&serde_json::json!({ "name": name })).unwrap(), + pub fn new(name: Option<&str>) -> Self { + let body: Vec = match name { + Some(name) => serde_json::to_vec(&serde_json::json!({ "name": name })) + .expect("json object serialization is infallible"), None => b"{}".to_vec(), }; Self { mount: "cert".to_owned(), - name, body, } } diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index 7937fa93f8f..3ad9a549438 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -207,7 +207,7 @@ impl HashicorpVault { (None, Some(tls)) => { let login_client: VaultClient = self.build_client(self.config.login_namespace(), "")?; - let endpoint: CertLoginRequest = CertLoginRequest::new(tls.role.clone()); + let endpoint: CertLoginRequest = CertLoginRequest::new(tls.role.as_deref()); let auth = api::auth(&login_client, endpoint) .await .map_err(|error| map_api_error(error, ErrorContext::Login))?; From 5c0589207cb1c78600b0407d003f0e1bc3822be7 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 22:30:13 +0000 Subject: [PATCH 11/14] fix(xai): accept max_completion_tokens as a supported param Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/xai/chat/transformation.py | 1 + .../llms/xai/test_xai_chat_transformation.py | 10 ++++++++++ 2 files changed, 11 insertions(+) diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 91bf697487d..33ee727dfab 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -118,6 +118,7 @@ class XAIChatConfig(OpenAIGPTConfig): base_openai_params: Final = [ "logit_bias", "logprobs", + "max_completion_tokens", "max_tokens", "n", "parallel_tool_calls", diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index 290cd3dcb3a..3fd666e4f50 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -90,6 +90,16 @@ class TestXAIReasoningTokenFolding: assert response.usage.total_tokens == 999 +def test_max_completion_tokens_is_accepted_and_mapped_to_max_tokens() -> None: + optional_params = litellm.get_optional_params( + model="grok-4.20", + custom_llm_provider="xai", + max_completion_tokens=64, + ) + assert optional_params["max_tokens"] == 64, optional_params + assert "max_completion_tokens" not in optional_params, optional_params + + class TestXAIParallelToolCalls: """Test suite for XAI parallel tool calls functionality.""" From 4ff251e6b1ae6e2bf89424af8c06a453ded64e9b Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 22:37:13 +0000 Subject: [PATCH 12/14] fix(xai): accept max_completion_tokens as a supported param Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/xai/test_xai_chat_transformation.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index 3fd666e4f50..0a9e3234f12 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -13,6 +13,7 @@ from litellm.types.utils import ( ModelResponse, Usage, ) +from litellm.utils import get_optional_params class TestXAIReasoningTokenFolding: @@ -309,3 +310,15 @@ class TestXAIReportedCost: 0.0, 0.0037756, ) + + +class TestXAIMaxCompletionTokens: + def test_max_completion_tokens_is_accepted_and_mapped_to_max_tokens(self): + result = get_optional_params( + model="grok-4.20-beta", + custom_llm_provider="xai", + max_completion_tokens=64, + ) + + assert result["max_tokens"] == 64 + assert "max_completion_tokens" not in result From 32c63e332eaca89c0c740d097cc7b7727a3d628c Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 22:38:29 +0000 Subject: [PATCH 13/14] Revert "fix(xai): accept max_completion_tokens as a supported param" This reverts commit 4ff251e6b1ae6e2bf89424af8c06a453ded64e9b. --- .../llms/xai/test_xai_chat_transformation.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index 0a9e3234f12..3fd666e4f50 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -13,7 +13,6 @@ from litellm.types.utils import ( ModelResponse, Usage, ) -from litellm.utils import get_optional_params class TestXAIReasoningTokenFolding: @@ -310,15 +309,3 @@ class TestXAIReportedCost: 0.0, 0.0037756, ) - - -class TestXAIMaxCompletionTokens: - def test_max_completion_tokens_is_accepted_and_mapped_to_max_tokens(self): - result = get_optional_params( - model="grok-4.20-beta", - custom_llm_provider="xai", - max_completion_tokens=64, - ) - - assert result["max_tokens"] == 64 - assert "max_completion_tokens" not in result From 636a60d3bbc9fbfc59fa594b128520eb3e43f433 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:51:14 +0000 Subject: [PATCH 14/14] chore(rust): update merged workspace lockfile Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 108 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 106 insertions(+), 2 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 1a4eb51af08..5acfa22e824 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2850,6 +2850,7 @@ dependencies = [ "litellm-secrets-azure", "litellm-secrets-cyberark", "litellm-secrets-google", + "litellm-secrets-hashicorp", "litellm-secrets-types", "moka", "reqwest 0.12.28", @@ -2947,6 +2948,26 @@ dependencies = [ "wiremock", ] +[[package]] +name = "litellm-secrets-hashicorp" +version = "0.1.0" +dependencies = [ + "litellm-core-utils", + "litellm-secrets-types", + "moka", + "rstest", + "rustify", + "rustify_derive", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.19", + "tokio", + "vaultrs", + "veil", + "wiremock", +] + [[package]] name = "litellm-secrets-types" version = "0.1.0" @@ -4185,6 +4206,40 @@ dependencies = [ "semver", ] +[[package]] +name = "rustify" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4800ce4c1cc2fec12c559dae2ddbf0e17fcee7569b796e6d75898efef443368b" +dependencies = [ + "anyhow", + "async-trait", + "bytes", + "http 1.4.2", + "reqwest 0.13.5", + "rustify_derive", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror 1.0.69", + "tracing", + "url", +] + +[[package]] +name = "rustify_derive" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78ea7fda74240f7410d0198b603a8a2f662acc7d76b6667a49f9b162cd8d9b4f" +dependencies = [ + "proc-macro2", + "quote", + "regex", + "serde_urlencoded", + "syn 1.0.109", + "synstructure 0.12.6", +] + [[package]] name = "rustix" version = "1.1.5" @@ -4737,6 +4792,17 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "2.0.119" @@ -4768,6 +4834,18 @@ dependencies = [ "futures-core", ] +[[package]] +name = "synstructure" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f36bdaa60a83aca3921b5259d5400cbf5e90fc51931376a9bd4a0eb79aa7210f" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", + "unicode-xid", +] + [[package]] name = "synstructure" version = "0.13.2" @@ -5182,6 +5260,7 @@ version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ + "log", "pin-project-lite", "tracing-attributes", "tracing-core", @@ -5366,6 +5445,12 @@ version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + [[package]] name = "unicode_categories" version = "0.1.1" @@ -5425,6 +5510,25 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "vaultrs" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30ffcc0e81025065dda612ec1e26a3d81bb16ef3062354873d17a35965d68522" +dependencies = [ + "async-trait", + "derive_builder", + "http 1.4.2", + "reqwest 0.13.5", + "rustify", + "rustify_derive", + "serde", + "serde_json", + "thiserror 2.0.19", + "tracing", + "url", +] + [[package]] name = "vcpkg" version = "0.2.15" @@ -5874,7 +5978,7 @@ dependencies = [ "proc-macro2", "quote", "syn 2.0.119", - "synstructure", + "synstructure 0.13.2", ] [[package]] @@ -5915,7 +6019,7 @@ dependencies = [ "proc-macro2", "quote", "syn 2.0.119", - "synstructure", + "synstructure 0.13.2", ] [[package]]