From b462be816bcc1619502fce03ace66cf3746557b3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 22:06:56 -0700 Subject: [PATCH] fix(credentials): publish cross-pod config sync on credential writes --- .../proxy/credential_endpoints/endpoints.py | 7 ++ .../credential_endpoints/test_endpoints.py | 115 ++++++++++++++++-- 2 files changed, 111 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index dc193cb8523..14f240b91df 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -15,6 +15,10 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.config_sync_pubsub import ( + coordination_redis_cache, + publish_config_change, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.utils import handle_exception_on_proxy, jsonify_object from litellm.repositories.credentials_repository import CredentialsRepository @@ -104,6 +108,7 @@ async def create_credential( ## ADD TO LITELLM ## CredentialAccessor.upsert_credentials([processed_credential]) + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_credentialstable") return {"success": True, "message": "Credential created successfully"} except Exception as e: @@ -243,6 +248,7 @@ async def delete_credential( ## DELETE FROM LITELLM ## litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name] + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_credentialstable") return {"success": True, "message": "Credential deleted successfully"} except Exception as e: return handle_exception_on_proxy(e) @@ -351,6 +357,7 @@ async def update_credential( litellm.credential_list = [c for c in litellm.credential_list if c.credential_name != credential_name] CredentialAccessor.upsert_credentials([updated_in_memory]) + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_credentialstable") return {"success": True, "message": "Credential updated successfully"} except Exception as e: raise handle_exception_on_proxy(e) diff --git a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py index dcd8e6881bd..dd725a70fe1 100644 --- a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py @@ -2,10 +2,8 @@ from unittest.mock import AsyncMock, MagicMock, patch -import pytest from fastapi.testclient import TestClient - from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app @@ -19,12 +17,17 @@ def _as_admin(): def _patch_credential(name: str, body: dict): + return _request_as_admin("PATCH", f"/credentials/{name}", body) + + +def _request_as_admin(method: str, path: str, body: dict | None): missing = object() previous_override = app.dependency_overrides.get(user_api_key_auth, missing) app.dependency_overrides[user_api_key_auth] = _as_admin try: - return client.patch( - f"/credentials/{name}", + return client.request( + method, + path, json=body, headers={"Authorization": "Bearer test-key"}, ) @@ -40,14 +43,19 @@ def test_update_credential_answers_404_when_the_credential_does_not_exist(): the exception the response body and lets FastAPI answer 200, so a write the handler rejected read as a success to every caller that checks the status. The dashboard's API client branches on the status, so it reported a failed edit as applied.""" - with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( - "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" - ) as repository: + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository, + ): repository.return_value.find_by_name = AsyncMock(return_value=None) response = _patch_credential( "definitely-not-there", - {"credential_name": "definitely-not-there", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, + { + "credential_name": "definitely-not-there", + "credential_values": {"api_key": "sk-x"}, + "credential_info": {}, + }, ) assert response.status_code == 404, f"rejected write answered {response.status_code}: {response.text}" @@ -73,9 +81,11 @@ def test_update_credential_still_answers_200_on_a_successful_write(): credential_values={"api_key": "sk-old"}, credential_info={"custom_llm_provider": "openai"}, ) - with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( - "litellm.proxy.proxy_server.master_key", "sk-test-master" - ), patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository: + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "sk-test-master"), + patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository, + ): repository.return_value.find_by_name = AsyncMock(return_value=stored) repository.return_value.update_by_name = AsyncMock(return_value=None) @@ -86,3 +96,86 @@ def test_update_credential_still_answers_200_on_a_successful_write(): assert response.status_code == 200, response.text assert response.json()["success"] is True + + +def test_create_credential_publishes_a_cross_pod_sync_event(): + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "sk-test-master"), + patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository, + patch("litellm.proxy.credential_endpoints.endpoints.publish_config_change", new_callable=AsyncMock) as publish, + ): + repository.return_value.create = AsyncMock(return_value=None) + + response = _request_as_admin( + "POST", + "/credentials", + { + "credential_name": "sync-create", + "credential_values": {"api_key": "sk-x", "api_base": "http://127.0.0.1:1"}, + "credential_info": {"custom_llm_provider": "openai"}, + }, + ) + + assert response.status_code == 200, response.text + publish.assert_awaited_once() + assert publish.await_args.kwargs["object_type"] == "litellm_credentialstable" + + +def test_update_credential_publishes_a_cross_pod_sync_event(): + stored = CredentialItem( + credential_name="sync-update", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "openai"}, + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "sk-test-master"), + patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository, + patch("litellm.proxy.credential_endpoints.endpoints.publish_config_change", new_callable=AsyncMock) as publish, + ): + repository.return_value.find_by_name = AsyncMock(return_value=stored) + repository.return_value.update_by_name = AsyncMock(return_value=None) + + response = _request_as_admin( + "PATCH", + "/credentials/sync-update", + {"credential_name": "sync-update", "credential_values": {"api_key": "sk-new"}, "credential_info": {}}, + ) + + assert response.status_code == 200, response.text + publish.assert_awaited_once() + assert publish.await_args.kwargs["object_type"] == "litellm_credentialstable" + + +def test_delete_credential_publishes_a_cross_pod_sync_event(): + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository, + patch("litellm.proxy.credential_endpoints.endpoints.publish_config_change", new_callable=AsyncMock) as publish, + ): + repository.return_value.delete_by_name = AsyncMock(return_value=None) + + response = _request_as_admin("DELETE", "/credentials/sync-delete", None) + + assert response.status_code == 200, response.text + publish.assert_awaited_once() + assert publish.await_args.kwargs["object_type"] == "litellm_credentialstable" + + +def test_rejected_update_does_not_publish_a_sync_event(): + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository, + patch("litellm.proxy.credential_endpoints.endpoints.publish_config_change", new_callable=AsyncMock) as publish, + ): + repository.return_value.find_by_name = AsyncMock(return_value=None) + + response = _request_as_admin( + "PATCH", + "/credentials/not-there", + {"credential_name": "not-there", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, + ) + + assert response.status_code == 404, response.text + publish.assert_not_awaited()