From b7b22d559e625fa2ecd32b09bdea31ed2b095046 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Wed, 17 Dec 2025 17:46:24 +0900 Subject: [PATCH] feat: allow per-team Vault overrides when storing keys --- .../docs/secret_managers/hashicorp_vault.md | 21 +++ .../proxy/hooks/key_management_event_hooks.py | 80 ++++++++++- .../hashicorp_secret_manager.py | 131 ++++++++++++++---- tests/litellm_utils_tests/test_hashicorp.py | 98 ++++++++++++- .../hooks/test_key_management_event_hooks.py | 3 +- 5 files changed, 291 insertions(+), 42 deletions(-) diff --git a/docs/my-website/docs/secret_managers/hashicorp_vault.md b/docs/my-website/docs/secret_managers/hashicorp_vault.md index 09619609cb..c5082ce7fb 100644 --- a/docs/my-website/docs/secret_managers/hashicorp_vault.md +++ b/docs/my-website/docs/secret_managers/hashicorp_vault.md @@ -197,3 +197,24 @@ When a Virtual Key is Created / Deleted on LiteLLM, LiteLLM will automatically c LiteLLM stores secret under the `prefix_for_stored_virtual_keys` path (default: `litellm/`) + +### Team-specific overrides (proxy) + +When running the LiteLLM proxy you can override the Vault location per team. Set +`Secret Manager Settings` on the team with the following structure: + +```json +{ + "namespace": "teams/team-a", + "mount": "kv-prod", + "path_prefix": "virtual-keys", + "data": "password" +} +``` + +- `namespace` – overrides the `X-Vault-Namespace` header. +- `mount` – which KV engine mount to use (defaults to `secret`). +- `path_prefix` – additional path segments between the mount and the secret name. +- `data` – the field name inside the KV payload (defaults to `key`). + +Whenever LiteLLM stores or deletes virtual keys for that team, these overrides are applied so you can keep each team’s credentials in its own namespace, mount, or field layout without changing the global Vault configuration. diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 3aa62eeeed..3213e70027 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -1,7 +1,7 @@ import asyncio import json from datetime import datetime, timezone -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional import litellm from litellm._logging import verbose_proxy_logger @@ -78,6 +78,7 @@ class KeyManagementEventHooks: await KeyManagementEventHooks._store_virtual_key_in_secret_manager( secret_name=data.key_alias or f"virtual-key-{response.token_id}", secret_token=response.key, + team_id=data.team_id, ) except Exception as e: verbose_proxy_logger.warning( @@ -150,7 +151,8 @@ class KeyManagementEventHooks: ) await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager( current_secret_name=initial_secret_name, - new_secret_name=data.key_alias or f"virtual-key-{response.token_id}", + new_secret_name=data.key_alias + or f"virtual-key-{response.token_id}", new_secret_value=response.key, ) except Exception as e: @@ -241,7 +243,9 @@ class KeyManagementEventHooks: pass @staticmethod - async def _store_virtual_key_in_secret_manager(secret_name: str, secret_token: str): + async def _store_virtual_key_in_secret_manager( + secret_name: str, secret_token: str, team_id: Optional[str] = None + ): """ Store a virtual key in the secret manager @@ -261,6 +265,9 @@ class KeyManagementEventHooks: description = getattr( litellm._key_management_settings, "description", None ) + optional_params = await KeyManagementEventHooks._get_secret_manager_optional_params( + team_id + ) verbose_proxy_logger.debug( f"Creating secret with {secret_name} and tags={tags} and description={description}" ) @@ -271,7 +278,8 @@ class KeyManagementEventHooks: ), description=description, secret_value=secret_token, - tags=tags + tags=tags, + optional_params=optional_params, ) @staticmethod @@ -329,18 +337,76 @@ class KeyManagementEventHooks: ) if isinstance(litellm.secret_manager_client, BaseSecretManager): + team_settings_cache: Dict[Optional[str], Optional[dict]] = {} for key in keys_being_deleted: if key.key_alias is not None: + team_id = getattr(key, "team_id", None) + if team_id not in team_settings_cache: + team_settings_cache[ + team_id + ] = await KeyManagementEventHooks._get_secret_manager_optional_params( + team_id + ) + optional_params = team_settings_cache[team_id] await litellm.secret_manager_client.async_delete_secret( secret_name=KeyManagementEventHooks._get_secret_name( key.key_alias - ) + ), + optional_params=optional_params, ) else: verbose_proxy_logger.warning( f"KeyManagementEventHooks._delete_virtual_key_from_secret_manager: Key alias not found for key {key.token}. Skipping deletion from secret manager." ) + @staticmethod + async def _get_secret_manager_optional_params( + team_id: Optional[str], + ) -> Optional[dict]: + if team_id is None: + return None + + try: + from litellm.proxy import proxy_server as proxy_server_module + except ImportError: + return None + + prisma_client = getattr(proxy_server_module, "prisma_client", None) + user_api_key_cache = getattr(proxy_server_module, "user_api_key_cache", None) + + if prisma_client is None or user_api_key_cache is None: + return None + + try: + from litellm.proxy.auth.auth_checks import get_team_object + + team_obj = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except Exception as exc: # pragma: no cover - defensive logging + verbose_proxy_logger.debug( + f"Unable to load team metadata for team_id={team_id}: {exc}" + ) + return None + + metadata = getattr(team_obj, "metadata", None) + if metadata is None: + return None + + if hasattr(metadata, "model_dump"): + metadata = metadata.model_dump() + + if not isinstance(metadata, dict): + return None + + team_settings = metadata.get("secret_manager_settings") + if isinstance(team_settings, dict) and team_settings: + return dict(team_settings) + + return None + @staticmethod def _is_email_sending_enabled() -> bool: """ @@ -453,7 +519,9 @@ class KeyManagementEventHooks: ) @staticmethod - async def _send_key_rotated_email(response: dict, existing_key_alias: Optional[str]): + async def _send_key_rotated_email( + response: dict, existing_key_alias: Optional[str] + ): """ Send key rotated email if email sending is enabled. diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index fe26f5c332..cad9ccc7a9 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -66,11 +66,13 @@ class HashicorpSecretManager(BaseSecretManager): def _verify_required_credentials_exist(self) -> None: """ Validate that at least one authentication method is configured. - + Raises: ValueError: If no valid authentication credentials are provided """ - if not self.vault_token and not (self.approle_role_id and self.approle_secret_id): + if not self.vault_token and not ( + self.approle_role_id and self.approle_secret_id + ): raise ValueError( "Missing Vault authentication credentials. Please set either:\n" " - HCP_VAULT_TOKEN for token-based auth, or\n" @@ -107,20 +109,20 @@ class HashicorpSecretManager(BaseSecretManager): ``` """ verbose_logger.debug("Using AppRole auth for Hashicorp Vault") - + # Check cache first cached_token = self.cache.get_cache(key="hcp_vault_approle_token") if cached_token: verbose_logger.debug("Using cached Vault token from AppRole auth") return cached_token - + # Vault endpoint for AppRole login login_url = f"{self.vault_addr}/v1/auth/{self.approle_mount_path}/login" headers = {} if hasattr(self, "vault_namespace") and self.vault_namespace: headers["X-Vault-Namespace"] = self.vault_namespace - + try: client = _get_httpx_client() resp = client.post( @@ -132,15 +134,15 @@ class HashicorpSecretManager(BaseSecretManager): }, ) resp.raise_for_status() - + auth_data = resp.json()["auth"] token = auth_data["client_token"] _lease_duration = auth_data["lease_duration"] - + verbose_logger.debug( f"Successfully obtained Vault token via AppRole auth. Lease duration: {_lease_duration}s" ) - + # Cache the token with its lease duration self.cache.set_cache( key="hcp_vault_approle_token", value=token, ttl=_lease_duration @@ -209,31 +211,102 @@ class HashicorpSecretManager(BaseSecretManager): def _get_tls_cert_auth_body(self) -> dict: return {"name": self.vault_cert_role} - def get_url(self, secret_name: str) -> str: + def get_url( + self, + secret_name: str, + namespace: Optional[str] = None, + mount_name: Optional[str] = None, + path_prefix: Optional[str] = None, + ) -> str: """ Constructs the Vault URL for KV v2 secrets. - + Format: {VAULT_ADDR}/v1/{NAMESPACE}/{MOUNT_NAME}/data/{PATH_PREFIX}/{SECRET_NAME} - + Examples: - Default: http://127.0.0.1:8200/v1/secret/data/mykey - With namespace: http://127.0.0.1:8200/v1/mynamespace/secret/data/mykey - With custom mount: http://127.0.0.1:8200/v1/kv/data/mykey - With path prefix: http://127.0.0.1:8200/v1/secret/data/myapp/mykey """ + resolved_namespace = self._sanitize_path_component( + namespace if namespace is not None else self.vault_namespace + ) + resolved_mount = self._sanitize_path_component( + mount_name if mount_name is not None else self.vault_mount_name + ) + if resolved_mount is None: + resolved_mount = "secret" + resolved_path_prefix = self._sanitize_path_component( + path_prefix if path_prefix is not None else self.vault_path_prefix + ) + _url = f"{self.vault_addr}/v1/" - if self.vault_namespace: - _url += f"{self.vault_namespace}/" - _url += f"{self.vault_mount_name}/data/" - if self.vault_path_prefix: - _url += f"{self.vault_path_prefix}/" + if resolved_namespace: + _url += f"{resolved_namespace}/" + _url += f"{resolved_mount}/data/" + if resolved_path_prefix: + _url += f"{resolved_path_prefix}/" _url += secret_name return _url + def _sanitize_plain_value(self, value: Optional[Union[str, int]]) -> Optional[str]: + if value is None: + return None + value_str = str(value).strip() + if value_str == "": + return None + return value_str + + def _sanitize_path_component( + self, value: Optional[Union[str, int]] + ) -> Optional[str]: + sanitized_value = self._sanitize_plain_value(value) + if sanitized_value is None: + return None + sanitized_value = sanitized_value.strip("/") + return sanitized_value or None + + def _extract_secret_manager_settings( + self, optional_params: Optional[dict] + ) -> Dict[str, Any]: + if not isinstance(optional_params, dict): + return {} + + candidate = optional_params.get("secret_manager_settings") + source = candidate if isinstance(candidate, dict) else optional_params + allowed_keys = {"namespace", "mount", "path_prefix", "data"} + return {k: source[k] for k in allowed_keys if k in source} + + def _build_secret_target( + self, secret_name: str, optional_params: Optional[dict] + ) -> Dict[str, Any]: + settings = self._extract_secret_manager_settings(optional_params) + + namespace = settings.get("namespace", self.vault_namespace) + mount = settings.get("mount", self.vault_mount_name) + path_prefix = settings.get("path_prefix", self.vault_path_prefix) + data_key_override = settings.get("data") + + data_key = self._sanitize_plain_value(data_key_override) or "key" + + url = self.get_url( + secret_name=secret_name, + namespace=namespace, + mount_name=mount, + path_prefix=path_prefix, + ) + + return { + "url": url, + "data_key": data_key, + "secret_name": secret_name, + } + def _get_request_headers(self) -> dict: """ Get the headers for Vault API requests. - + Authentication priority: 1. AppRole (if role_id and secret_id are configured) 2. TLS Certificate (if cert paths are configured) @@ -242,11 +315,11 @@ class HashicorpSecretManager(BaseSecretManager): # Priority 1: AppRole auth if self.approle_role_id and self.approle_secret_id: return {"X-Vault-Token": self._auth_via_approle()} - + # Priority 2: TLS cert auth if self.tls_cert_path and self.tls_key_path: return {"X-Vault-Token": self._auth_via_tls_cert()} - + # Priority 3: Direct token return {"X-Vault-Token": self.vault_token} @@ -323,7 +396,7 @@ class HashicorpSecretManager(BaseSecretManager): description: Optional[str] = None, optional_params: Optional[dict] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - tags: Optional[Union[dict, list]] = None + tags: Optional[Union[dict, list]] = None, ) -> Dict[str, Any]: """ Writes a secret to Vault KV v2 using an async HTTPX client. @@ -344,16 +417,18 @@ class HashicorpSecretManager(BaseSecretManager): ) try: - url = self.get_url(secret_name) + target = self._build_secret_target(secret_name, optional_params) # Prepare the secret data - data = {"data": {"key": secret_value}} + data = {"data": {target["data_key"]: secret_value}} if description: data["data"]["description"] = description response = await async_client.post( - url=url, headers=self._get_request_headers(), json=data + url=target["url"], + headers=self._get_request_headers(), + json=data, ) response.raise_for_status() return response.json() @@ -397,20 +472,20 @@ class HashicorpSecretManager(BaseSecretManager): ) try: - # For KV v2 delete: /v1//data/ - url = self.get_url(secret_name) - + target = self._build_secret_target(secret_name, optional_params) response = await async_client.delete( - url=url, headers=self._get_request_headers() + url=target["url"], headers=self._get_request_headers() ) response.raise_for_status() # Clear the cache for this secret self.cache.delete_cache(secret_name) + if target["secret_name"] != secret_name: + self.cache.delete_cache(target["secret_name"]) return { "status": "success", - "message": f"Secret {secret_name} deleted successfully", + "message": f"Secret {target['secret_name']} deleted successfully", } except Exception as e: verbose_logger.exception(f"Error deleting secret from Hashicorp Vault: {e}") diff --git a/tests/litellm_utils_tests/test_hashicorp.py b/tests/litellm_utils_tests/test_hashicorp.py index 4757c72262..4f3536f9bf 100644 --- a/tests/litellm_utils_tests/test_hashicorp.py +++ b/tests/litellm_utils_tests/test_hashicorp.py @@ -23,7 +23,16 @@ litellm.proxy.proxy_server.premium_user = True from litellm.secret_managers.hashicorp_secret_manager import HashicorpSecretManager -hashicorp_secret_manager = HashicorpSecretManager() + +@pytest.fixture +def hashicorp_secret_manager(): + """Provide a fresh HashicorpSecretManager per test to avoid shared state.""" + manager = HashicorpSecretManager() + manager.vault_addr = "https://test-cluster-public-vault-0f98180c.e98296b2.z1.hashicorp.cloud:8200" + manager.vault_namespace = "admin" + manager.vault_mount_name = "secret" + manager.vault_path_prefix = None + return manager mock_vault_response = { @@ -67,7 +76,7 @@ mock_write_response = { } -def test_hashicorp_secret_manager_get_secret(): +def test_hashicorp_secret_manager_get_secret(hashicorp_secret_manager): with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get") as mock_get: # Configure the mock response using MagicMock mock_response = MagicMock() @@ -92,7 +101,7 @@ def test_hashicorp_secret_manager_get_secret(): @pytest.mark.asyncio -async def test_hashicorp_secret_manager_write_secret(): +async def test_hashicorp_secret_manager_write_secret(hashicorp_secret_manager): with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" ) as mock_post: @@ -136,7 +145,47 @@ async def test_hashicorp_secret_manager_write_secret(): @pytest.mark.asyncio -async def test_hashicorp_secret_manager_delete_secret(): +async def test_hashicorp_secret_manager_write_secret_with_team_overrides( + hashicorp_secret_manager, +): + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" + ) as mock_post: + mock_response = MagicMock() + mock_response.json.return_value = mock_write_response + mock_response.raise_for_status.return_value = None + mock_post.return_value = mock_response + + secret_value = "value-mock" + team_settings = { + "namespace": "team-namespace", + "mount": "kv-team", + "path_prefix": "teams/custom", + "data": "password", + } + + response = await hashicorp_secret_manager.async_write_secret( + secret_name="team-secret", + secret_value=secret_value, + optional_params=team_settings, + ) + + assert response == mock_write_response + mock_post.assert_called_once() + + called_url = mock_post.call_args[1]["url"] + expected_url = ( + f"{hashicorp_secret_manager.vault_addr}/v1/" + "team-namespace/kv-team/data/teams/custom/team-secret" + ) + assert called_url == expected_url + + json_data = mock_post.call_args[1]["json"] + assert json_data["data"] == {"password": secret_value} + + +@pytest.mark.asyncio +async def test_hashicorp_secret_manager_delete_secret(hashicorp_secret_manager): with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.delete" ) as mock_delete: @@ -169,7 +218,42 @@ async def test_hashicorp_secret_manager_delete_secret(): ) -def test_hashicorp_secret_manager_tls_cert_auth(monkeypatch): +@pytest.mark.asyncio +async def test_hashicorp_secret_manager_delete_secret_with_team_overrides( + hashicorp_secret_manager, +): + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.delete" + ) as mock_delete: + mock_response = MagicMock() + mock_response.raise_for_status.return_value = None + mock_delete.return_value = mock_response + + team_settings = { + "namespace": "team-namespace", + "mount": "kv-team", + "path_prefix": "teams/custom", + } + + response = await hashicorp_secret_manager.async_delete_secret( + secret_name="team-secret", optional_params=team_settings + ) + + assert response == { + "status": "success", + "message": "Secret team-secret deleted successfully", + } + + mock_delete.assert_called_once() + called_url = mock_delete.call_args[1]["url"] + expected_url = ( + f"{hashicorp_secret_manager.vault_addr}/v1/" + "team-namespace/kv-team/data/teams/custom/team-secret" + ) + assert called_url == expected_url + + +def test_hashicorp_secret_manager_tls_cert_auth(monkeypatch, hashicorp_secret_manager): monkeypatch.setenv("HCP_VAULT_TOKEN", "test-client-token-12345") print("HCP_VAULT_TOKEN=", os.getenv("HCP_VAULT_TOKEN")) # Mock both httpx.post and httpx.Client @@ -217,7 +301,7 @@ def test_hashicorp_secret_manager_tls_cert_auth(monkeypatch): assert test_manager.cache.get_cache("hcp_vault_token") == "test-client-token-12345" -def test_hashicorp_secret_manager_approle_auth(monkeypatch): +def test_hashicorp_secret_manager_approle_auth(monkeypatch, hashicorp_secret_manager): """ Test AppRole authentication makes the expected POST request to the correct URL. """ @@ -260,7 +344,7 @@ def test_hashicorp_secret_manager_approle_auth(monkeypatch): assert test_manager.cache.get_cache("hcp_vault_approle_token") == "hvs.approle-token-67890" -def test_hashicorp_custom_mount_and_prefix(): +def test_hashicorp_custom_mount_and_prefix(hashicorp_secret_manager): """Test URL construction with custom mount name and path prefix using get_url method.""" # Save original values original_mount = hashicorp_secret_manager.vault_mount_name diff --git a/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py index f731d9e298..011031c1e4 100644 --- a/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py +++ b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py @@ -39,6 +39,7 @@ class TestKeyManagementEventHooksIndependentOperations: # Create mock objects for the hook parameters mock_data = MagicMock() mock_data.key_alias = "test-key-alias" + mock_data.team_id = None mock_response = MagicMock() mock_response.model_dump.return_value = {"key": "sk-test", "token": "test-token"} @@ -94,6 +95,7 @@ class TestKeyManagementEventHooksIndependentOperations: # Create mock objects for the hook parameters mock_data = MagicMock() mock_data.key_alias = "test-key-alias" + mock_data.team_id = None mock_response = MagicMock() mock_response.model_dump.return_value = {"key": "sk-test", "token": "test-token"} @@ -127,4 +129,3 @@ class TestKeyManagementEventHooksIndependentOperations: # Email should have been called despite secret manager failure assert email_called["called"] is True -