From 7b78e1426b426dc90516d5040a8d2733ade04d47 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 21 Jul 2026 06:11:41 +0000 Subject: [PATCH] fix(proxy): harden credential scope and migration recovery Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_proxy_extras/utils.py | 190 ++++++------------ .../tests/test_setup_database_fail_fast.py | 73 ++++--- .../proxy/credential_endpoints/endpoints.py | 51 ++--- .../credential_endpoints/test_endpoints.py | 146 +++++++------- 4 files changed, 196 insertions(+), 264 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 369b6561931..f7c0bed0800 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -24,9 +24,7 @@ def _get_prisma_env() -> dict: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # These env vars prevent Prisma from attempting downloads prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true" - prisma_env["NPM_CONFIG_CACHE"] = os.getenv( - "NPM_CONFIG_CACHE", "/app/.cache/npm" - ) + prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm") return prisma_env @@ -69,8 +67,7 @@ def _get_prisma_command() -> str: # If not found, log warning and fall back logger.warning( - f"Prisma CLI not found at {default_cli_path}. " - "Falling back to Python wrapper (may attempt downloads)" + f"Prisma CLI not found at {default_cli_path}. Falling back to Python wrapper (may attempt downloads)" ) # Fall back to the Python wrapper (will work in online mode) @@ -159,14 +156,10 @@ class ProxyExtrasDBManager: return True except subprocess.TimeoutExpired: - logger.warning( - "Migration timed out - the database might be under heavy load." - ) + logger.warning("Migration timed out - the database might be under heavy load.") return False except subprocess.CalledProcessError as e: - logger.warning( - f"Error creating baseline migration: {e}, {e.stderr}, {e.stdout}" - ) + logger.warning(f"Error creating baseline migration: {e}, {e.stderr}, {e.stdout}") raise e @staticmethod @@ -264,9 +257,7 @@ class ProxyExtrasDBManager: return False @staticmethod - def _resolve_all_migrations( - migrations_dir: str, schema_path: str, mark_all_applied: bool = True - ): + def _resolve_all_migrations(migrations_dir: str, schema_path: str, mark_all_applied: bool = True): """ 1. Compare the current database state to schema.prisma and generate a migration for the diff. 2. Run prisma migrate deploy to apply any pending migrations. @@ -339,9 +330,7 @@ class ProxyExtrasDBManager: ) logger.info(f"Applied migration: {mig_file.parent.name}") except subprocess.CalledProcessError as e: - logger.warning( - f"Failed to apply migration {mig_file.parent.name}: {e.stderr}" - ) + logger.warning(f"Failed to apply migration {mig_file.parent.name}: {e.stderr}") except subprocess.TimeoutExpired: logger.warning(f"Migration {mig_file.parent.name} timed out.") return @@ -398,9 +387,7 @@ class ProxyExtrasDBManager: logger.debug(f"Resolved migration: {migration_name}") except subprocess.CalledProcessError as e: if "is already recorded as applied in the database." not in e.stderr: - logger.warning( - f"Failed to resolve migration {migration_name}: {e.stderr}" - ) + logger.warning(f"Failed to resolve migration {migration_name}: {e.stderr}") @staticmethod def _strip_prisma_query_params(url: str) -> str: @@ -460,9 +447,7 @@ class ProxyExtrasDBManager: # it, psycopg3's `with conn` calls COMMIT on clean exit — which # fails after `UndefinedTable` (fresh DB) leaves the transaction # in an aborted state. - with psycopg.connect( - cleaned_url, connect_timeout=10, autocommit=True - ) as conn: + with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: try: rows = conn.execute( "SELECT migration_name FROM _prisma_migrations " @@ -483,9 +468,7 @@ class ProxyExtrasDBManager: return head_newest_ts = _max_migration_timestamp(known) - hostile = { - name for name in unknown if _migration_timestamp(name) > head_newest_ts - } + hostile = {name for name in unknown if _migration_timestamp(name) > head_newest_ts} if not hostile: return @@ -546,8 +529,11 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) + timeout_attempts = 0 + baseline_created = False + recovered_migrations: frozenset[str] = frozenset() try: - for attempt in range(4): + while True: try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], @@ -561,9 +547,12 @@ class ProxyExtrasDBManager: return True except subprocess.TimeoutExpired: - logger.info( - f"prisma migrate deploy attempt {attempt + 1} timed out, retrying" - ) + timeout_attempts += 1 + if timeout_attempts >= 4: + raise RuntimeError( + "Database migration failed after 4 timeout attempts. Check database connectivity and load." + ) + logger.info(f"prisma migrate deploy attempt {timeout_attempts} timed out, retrying") time.sleep(random.randrange(5, 15)) continue @@ -571,22 +560,25 @@ class ProxyExtrasDBManager: stderr = e.stderr or "" if "P3005" in stderr and "database schema is not empty" in stderr: - logger.info( - "Schema exists but no migrations ledger — creating baseline" - ) + if baseline_created: + raise RuntimeError( + "Database migration repeatedly reported a missing " + "migration ledger after baseline creation." + ) from e + logger.info("Schema exists but no migrations ledger — creating baseline") ProxyExtrasDBManager._create_baseline_migration(schema_path) + baseline_created = True continue if "P3009" in stderr: migration_match = re.search(r"`(\d+_\S+?)`", stderr) - if ( - migration_match - and ProxyExtrasDBManager._is_idempotent_error(stderr) - ): + if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr): name = migration_match.group(1) - logger.info( - f"Migration {name} failed idempotently — marking applied and retrying" - ) + if name in recovered_migrations: + raise RuntimeError( + f"Migration {name} remained failed after idempotent recovery." + ) from e + logger.info(f"Migration {name} failed idempotently — marking applied and retrying") try: ProxyExtrasDBManager._roll_back_migration(name) except ( @@ -611,6 +603,7 @@ class ProxyExtrasDBManager: f"intervention may be required.\n\n" f"Detail: {resolve_err}" ) from resolve_err + recovered_migrations = recovered_migrations | {name} continue raise RuntimeError( "Database migration failed and cannot be auto-recovered. " @@ -625,17 +618,14 @@ class ProxyExtrasDBManager: f"and retry.\n\nPrisma error:\n{stderr}" ) from e - migration_match = re.search( - r"Migration name: (\d+_\S+)", stderr - ) - if ( - migration_match - and ProxyExtrasDBManager._is_idempotent_error(stderr) - ): + migration_match = re.search(r"Migration name: (\d+_\S+)", stderr) + if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr): name = migration_match.group(1) - logger.info( - f"Migration {name} SQL hit idempotent error — marking applied and retrying" - ) + if name in recovered_migrations: + raise RuntimeError( + f"Migration {name} remained failed after idempotent recovery." + ) from e + logger.info(f"Migration {name} SQL hit idempotent error — marking applied and retrying") try: ProxyExtrasDBManager._roll_back_migration(name) except ( @@ -655,6 +645,7 @@ class ProxyExtrasDBManager: f"intervention may be required.\n\n" f"Detail: {resolve_err}" ) from resolve_err + recovered_migrations = recovered_migrations | {name} continue raise RuntimeError( @@ -667,19 +658,11 @@ class ProxyExtrasDBManager: f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e - raise RuntimeError( - "Database migration failed after 4 attempts (retry loop " - "exhausted by timeouts or repeated idempotent-recovery " - "continues). Check database connectivity, load, and " - "_prisma_migrations ledger state." - ) finally: os.chdir(original_dir) @staticmethod - def setup_database( - use_migrate: bool = False, use_v2_resolver: bool = False - ) -> bool: + def setup_database(use_migrate: bool = False, use_v2_resolver: bool = False) -> bool: """ Set up the database using either prisma migrate or prisma db push Uses migrations from litellm-proxy-extras package @@ -724,9 +707,7 @@ class ProxyExtrasDBManager: # Skip sanity check when deploy reports no pending migrations — # DB already matches schema, no drift to correct. if "No pending migrations to apply" in result.stdout: - logger.info( - "No pending migrations — skipping post-migration sanity check" - ) + logger.info("No pending migrations — skipping post-migration sanity check") return True # Run sanity check to ensure DB matches schema @@ -740,9 +721,7 @@ class ProxyExtrasDBManager: logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}") if "P3009" in e.stderr: # Extract the failed migration name from the error message - migration_match = re.search( - r"`(\d+_.*)` migration", e.stderr - ) + migration_match = re.search(r"`(\d+_.*)` migration", e.stderr) if migration_match: failed_migration = migration_match.group(1) if ProxyExtrasDBManager._is_idempotent_error(e.stderr): @@ -750,9 +729,7 @@ class ProxyExtrasDBManager: f"Migration {failed_migration} failed due to idempotent error (e.g., column already exists), resolving as applied" ) try: - ProxyExtrasDBManager._roll_back_migration( - failed_migration - ) + ProxyExtrasDBManager._roll_back_migration(failed_migration) except ( subprocess.CalledProcessError, subprocess.TimeoutExpired, @@ -762,9 +739,7 @@ class ProxyExtrasDBManager: f"It may already be in a rolled-back state." ) try: - ProxyExtrasDBManager._resolve_specific_migration( - failed_migration - ) + ProxyExtrasDBManager._resolve_specific_migration(failed_migration) logger.info( f"✅ Migration {failed_migration} resolved, retrying to apply remaining migrations" ) @@ -772,9 +747,7 @@ class ProxyExtrasDBManager: subprocess.CalledProcessError, subprocess.TimeoutExpired, ) as resolve_err: - logger.warning( - f"Failed to resolve migration {failed_migration}: {resolve_err}" - ) + logger.warning(f"Failed to resolve migration {failed_migration}: {resolve_err}") # Apply any schema drift not covered by the marked-as-applied migration ProxyExtrasDBManager._resolve_all_migrations( migrations_dir, @@ -782,9 +755,7 @@ class ProxyExtrasDBManager: mark_all_applied=False, ) else: - logger.info( - f"Found failed migration: {failed_migration}, marking as rolled back" - ) + logger.info(f"Found failed migration: {failed_migration}, marking as rolled back") # Mark the failed migration as rolled back subprocess.run( [ @@ -800,23 +771,14 @@ class ProxyExtrasDBManager: text=True, env=_get_prisma_env(), ) - logger.info( - f"✅ Migration {failed_migration} marked as rolled back... retrying" - ) - elif ( - "P3005" in e.stderr - and "database schema is not empty" in e.stderr - ): + logger.info(f"✅ Migration {failed_migration} marked as rolled back... retrying") + elif "P3005" in e.stderr and "database schema is not empty" in e.stderr: logger.info( "Database schema is not empty, creating baseline migration. In read-only file system, please set an environment variable `LITELLM_MIGRATION_DIR` to a writable directory to enable migrations. Learn more - https://docs.litellm.ai/docs/proxy/prod#read-only-file-system" ) ProxyExtrasDBManager._create_baseline_migration(schema_path) - logger.info( - "Baseline migration created, resolving all migrations" - ) - ProxyExtrasDBManager._resolve_all_migrations( - migrations_dir, schema_path - ) + logger.info("Baseline migration created, resolving all migrations") + ProxyExtrasDBManager._resolve_all_migrations(migrations_dir, schema_path) logger.info("✅ All migrations resolved.") return True elif "P3018" in e.stderr: @@ -824,14 +786,8 @@ class ProxyExtrasDBManager: if ProxyExtrasDBManager._is_permission_error(e.stderr): # Permission errors should NOT be marked as applied # Extract migration name for logging - migration_match = re.search( - r"Migration name: (\d+_.*)", e.stderr - ) - migration_name = ( - migration_match.group(1) - if migration_match - else "unknown" - ) + migration_match = re.search(r"Migration name: (\d+_.*)", e.stderr) + migration_name = migration_match.group(1) if migration_match else "unknown" logger.error( f"❌ Migration {migration_name} failed due to insufficient permissions. " @@ -841,16 +797,10 @@ class ProxyExtrasDBManager: # Mark as rolled back and exit with error if migration_match: try: - ProxyExtrasDBManager._roll_back_migration( - migration_name - ) - logger.info( - f"Migration {migration_name} marked as rolled back" - ) + ProxyExtrasDBManager._roll_back_migration(migration_name) + logger.info(f"Migration {migration_name} marked as rolled back") except Exception as rollback_error: - logger.warning( - f"Failed to mark migration as rolled back: {rollback_error}" - ) + logger.warning(f"Failed to mark migration as rolled back: {rollback_error}") # Re-raise the error to prevent silent failures raise RuntimeError( @@ -865,18 +815,12 @@ class ProxyExtrasDBManager: "resolving as applied" ) # Extract the migration name from the error message - migration_match = re.search( - r"Migration name: (\d+_.*)", e.stderr - ) + migration_match = re.search(r"Migration name: (\d+_.*)", e.stderr) if migration_match: migration_name = migration_match.group(1) try: - logger.info( - f"Rolling back migration {migration_name}" - ) - ProxyExtrasDBManager._roll_back_migration( - migration_name - ) + logger.info(f"Rolling back migration {migration_name}") + ProxyExtrasDBManager._roll_back_migration(migration_name) except ( subprocess.CalledProcessError, subprocess.TimeoutExpired, @@ -890,9 +834,7 @@ class ProxyExtrasDBManager: f"Resolving migration {migration_name} that failed " f"due to existing schema objects" ) - ProxyExtrasDBManager._resolve_specific_migration( - migration_name - ) + ProxyExtrasDBManager._resolve_specific_migration(migration_name) logger.info( f"✅ Migration {migration_name} resolved, " f"retrying to apply remaining migrations" @@ -901,9 +843,7 @@ class ProxyExtrasDBManager: subprocess.CalledProcessError, subprocess.TimeoutExpired, ) as resolve_err: - logger.warning( - f"Failed to resolve migration {migration_name}: {resolve_err}" - ) + logger.warning(f"Failed to resolve migration {migration_name}: {resolve_err}") # Apply any schema drift not covered by the marked-as-applied migration ProxyExtrasDBManager._resolve_all_migrations( migrations_dir, @@ -931,11 +871,7 @@ class ProxyExtrasDBManager: time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: attempts_left = 3 - attempt - retry_msg = ( - f" Retrying... ({attempts_left} attempts left)" - if attempts_left > 0 - else "" - ) + retry_msg = f" Retrying... ({attempts_left} attempts left)" if attempts_left > 0 else "" logger.info(f"The process failed to execute. Details: {e}.{retry_msg}") time.sleep(random.randrange(5, 15)) finally: diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 8d66bf872de..fb83872bf6c 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -6,7 +6,7 @@ The v2 resolver is opt-in via `--use_v2_migration_resolver` / the """ import subprocess -from unittest.mock import patch +from unittest.mock import MagicMock, call, patch import pytest @@ -32,9 +32,7 @@ def _fake_migrate_deploy_failure(returncode: int, stderr: str): def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): """v2: a permission failure during migrate deploy raises RuntimeError.""" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") @@ -50,9 +48,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): """v2: a non-idempotent migration failure raises (no silent recovery).""" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") @@ -176,20 +172,14 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) -def test_v2_resolve_specific_migration_failure_raises_runtime_error( - monkeypatch, tmp_path -): +def test_v2_resolve_specific_migration_failure_raises_runtime_error(monkeypatch, tmp_path): """If marking a migration as applied fails inside P3009 idempotent recovery, the subprocess error must be re-raised as RuntimeError so proxy_cli.py catches it cleanly (instead of leaking CalledProcessError).""" - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") - monkeypatch.setattr( - ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None) # First call: migrate deploy -> P3009 idempotent error. # Recovery path tries _resolve_specific_migration; that also raises. @@ -201,26 +191,17 @@ def test_v2_resolve_specific_migration_failure_raises_runtime_error( output="", ) - monkeypatch.setattr( - ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve) - stderr = ( - "Error: P3009\nMigration `20260101000000_some_migration` failed\n" - "relation already exists" - ) + stderr = "Error: P3009\nMigration `20260101000000_some_migration` failed\nrelation already exists" with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): - with pytest.raises( - RuntimeError, match="Failed to mark migration .* as applied" - ): + with pytest.raises(RuntimeError, match="Failed to mark migration .* as applied"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): """v2 must never call _resolve_all_migrations — that's the bug it fixes.""" - monkeypatch.setattr( - ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None - ) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) (tmp_path / "schema.prisma").write_text("// stub") @@ -240,3 +221,37 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery" + + +def test_v2_recovers_multiple_idempotent_migrations_in_one_invocation(monkeypatch, tmp_path): + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + migration_names = tuple(f"2026010100000{index}_migration_{index}" for index in range(5)) + failures = tuple( + subprocess.CalledProcessError( + returncode=1, + cmd="prisma migrate deploy", + stderr=(f"Error: P3018\nMigration name: {migration_name}\nrelation already exists"), + output="", + ) + for migration_name in migration_names + ) + + class FakeResult: + stdout = "Applied migration.\n" + stderr = "" + + run = MagicMock(side_effect=(*failures, FakeResult())) + roll_back = MagicMock() + resolve = MagicMock() + monkeypatch.setattr("subprocess.run", run) + monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", roll_back) + monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_specific_migration", resolve) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert ok is True + assert run.call_count == 6 + assert roll_back.call_args_list == [call(name) for name in migration_names] + assert resolve.call_args_list == [call(name) for name in migration_names] diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 165b1b9e666..0db1d6bddb4 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -160,34 +160,6 @@ async def _caller_grantable_team_ids( return (await _caller_admin_scope(user_api_key_dict, prisma_client)).team_ids -def _credential_in_memory(credential_name: str) -> CredentialItem | None: - return next( - (cred for cred in litellm.credential_list if cred.credential_name == credential_name), - None, - ) - - -async def _credential_for_admin_gate(credential_name: str, prisma_client: object) -> CredentialItem | None: - """Authoritative credential lookup for the admin gate on update/delete. - - The in-process ``litellm.credential_list`` can be stale: a credential created - via the API on another horizontally-scaled instance, or before a restart, - exists only in the DB. Gating on the in-memory copy alone would let a logging - credential that isn't resident be updated/deleted without the proxy-admin - check. Prefer the in-memory copy, fall back to the DB so the gate sees the - real ``credential_info``. - """ - existing = _credential_in_memory(credential_name) - if existing is not None: - return existing - if prisma_client is None: - return None - try: - return await CredentialsRepository(prisma_client).find_by_name(credential_name) - except Exception: # noqa: BLE001 # treat any lookup failure as credential-not-found - return None - - class CredentialHelperUtils: @staticmethod def encrypt_credential_values( @@ -615,7 +587,15 @@ async def update_credential( """ from litellm.proxy.proxy_server import prisma_client - existing = await _credential_for_admin_gate(credential_name, prisma_client) + if prisma_client is None: + return handle_exception_on_proxy( + HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + ) + credentials_repository = CredentialsRepository(prisma_client) + existing = await credentials_repository.find_by_name(credential_name) await _authorize_credential_patch( credential_name=credential_name, patch=credential, @@ -623,19 +603,12 @@ async def update_credential( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) + if existing is None: + return handle_exception_on_proxy(HTTPException(status_code=404, detail="Credential not found in DB.")) validate_credential_access(credential.credential_info) try: - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - credentials_repository = CredentialsRepository(prisma_client) - db_credential = await credentials_repository.find_by_name(credential_name) - if db_credential is None: - raise HTTPException(status_code=404, detail="Credential not found in DB.") - merged_credential = update_db_credential(db_credential, _patch_to_credential_item(credential, credential_name)) + merged_credential = update_db_credential(existing, _patch_to_credential_item(credential, credential_name)) credential_object_jsonified = jsonify_object(merged_credential.model_dump()) await credentials_repository.update_by_name( credential_name, diff --git a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py index 2874dedd3d4..5337ef7e1f8 100644 --- a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py @@ -48,10 +48,17 @@ def _connected_db(monkeypatch): repo = MagicMock() repo.create = AsyncMock() repo.delete_by_name = AsyncMock() + repo.update_by_name = AsyncMock() + + async def _find_by_name(name: str) -> CredentialItem | None: + return next( + (credential for credential in litellm.credential_list if credential.credential_name == name), + None, + ) + + repo.find_by_name = AsyncMock(side_effect=_find_by_name) monkeypatch.setattr(endpoints, "CredentialsRepository", lambda _client: repo) - monkeypatch.setattr( - endpoints.CredentialAccessor, "upsert_credentials", lambda creds: None - ) + monkeypatch.setattr(endpoints.CredentialAccessor, "upsert_credentials", lambda creds: None) return repo @@ -130,9 +137,7 @@ async def test_update_logging_credential_forbidden_for_non_admin(_connected_db): @pytest.mark.asyncio -async def test_update_existing_logging_credential_forbidden_even_without_logging_patch( - _connected_db, monkeypatch -): +async def test_update_existing_logging_credential_forbidden_even_without_logging_patch(_connected_db, monkeypatch): """A non-admin cannot edit a stored logging credential's values, even with a patch that omits credential_info (the gate consults the in-memory credential too).""" monkeypatch.setattr( @@ -234,9 +239,7 @@ def test_update_db_credential_preserves_untouched_access_subfields(): @pytest.mark.asyncio -async def test_delete_logging_credential_forbidden_for_non_admin( - _connected_db, monkeypatch -): +async def test_delete_logging_credential_forbidden_for_non_admin(_connected_db, monkeypatch): monkeypatch.setattr( litellm, "credential_list", @@ -260,9 +263,7 @@ async def test_delete_logging_credential_forbidden_for_non_admin( @pytest.mark.asyncio -async def test_update_db_only_logging_credential_forbidden_for_non_admin( - _connected_db, monkeypatch -): +async def test_update_db_only_logging_credential_forbidden_for_non_admin(_connected_db, monkeypatch): """A logging credential that exists ONLY in the DB (not resident in the in-memory ``credential_list`` -- e.g. created on another scaled instance or before a restart) must still gate a non-admin update. The gate falls back to @@ -292,9 +293,7 @@ async def test_update_db_only_logging_credential_forbidden_for_non_admin( @pytest.mark.asyncio -async def test_delete_db_only_logging_credential_forbidden_for_non_admin( - _connected_db, monkeypatch -): +async def test_delete_db_only_logging_credential_forbidden_for_non_admin(_connected_db, monkeypatch): """Same DB-only fallback for delete: a non-admin can't delete a logging credential that is resident only in the DB.""" monkeypatch.setattr(litellm, "credential_list", []) @@ -332,9 +331,7 @@ def _team_admin_of(team_ids): Combined with ``_patch_team_admin_lookup`` it mimics the real ``_caller_grantable_team_ids`` resolution without touching the DB. """ - return UserAPIKeyAuth( - api_key="k", user_role=LitellmUserRoles.INTERNAL_USER, user_id="ta-demo" - ) + return UserAPIKeyAuth(api_key="k", user_role=LitellmUserRoles.INTERNAL_USER, user_id="ta-demo") @pytest.fixture @@ -368,9 +365,7 @@ def _resident_logging_dest(): @pytest.mark.asyncio -async def test_team_admin_can_append_own_team_to_access( - _connected_db, _patch_team_admin_lookup, monkeypatch -): +async def test_team_admin_can_append_own_team_to_access(_connected_db, _patch_team_admin_lookup, monkeypatch): monkeypatch.setattr(litellm, "credential_list", [_resident_logging_dest()]) _connected_db.find_by_name = AsyncMock(return_value=_resident_logging_dest()) _connected_db.update_by_name = AsyncMock() @@ -388,13 +383,56 @@ async def test_team_admin_can_append_own_team_to_access( user_api_key_dict=_team_admin_of(["team-T"]), ) assert result["success"] is True + _connected_db.find_by_name.assert_awaited_once_with("dest") _connected_db.update_by_name.assert_awaited_once() @pytest.mark.asyncio -async def test_provider_credential_patch_forbidden_for_non_admin( - _connected_db, monkeypatch -): +async def test_team_admin_cannot_replay_stale_global_scope(_connected_db, _patch_team_admin_lookup, monkeypatch): + cached = CredentialItem( + credential_name="dest", + credential_values={"langfuse_host": "h"}, + credential_info={ + **_DEST_WITH_TEAMS, + "access": {"global": False, "teams": ["team-existing"]}, + }, + ) + authoritative = CredentialItem( + credential_name="dest", + credential_values={"langfuse_host": "h"}, + credential_info={ + **_DEST_WITH_TEAMS, + "access": {"global": True, "teams": ["team-existing"]}, + }, + ) + monkeypatch.setattr(litellm, "credential_list", [cached]) + _connected_db.find_by_name = AsyncMock(return_value=authoritative) + _connected_db.update_by_name = AsyncMock() + _patch_team_admin_lookup["ids"] = frozenset({"team-T"}) + + with pytest.raises(HTTPException) as exc: + await endpoints.update_credential( + request=MagicMock(), + fastapi_response=MagicMock(), + credential=UpdateCredentialItem( + credential_info={ + "access": { + "global": False, + "teams": ["team-existing", "team-T"], + } + }, + ), + credential_name="dest", + user_api_key_dict=_team_admin_of(["team-T"]), + ) + + assert exc.value.status_code == 403 + _connected_db.find_by_name.assert_awaited_once_with("dest") + _connected_db.update_by_name.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_provider_credential_patch_forbidden_for_non_admin(_connected_db, monkeypatch): """A team-admin (or any non-admin) cannot PATCH a non-logging credential. The route gate was widened to let team-admins reach /credentials/{name} @@ -428,9 +466,7 @@ async def test_provider_credential_patch_forbidden_for_non_admin( @pytest.mark.asyncio -async def test_provider_credential_access_patch_bypass_forbidden( - _connected_db, _patch_team_admin_lookup, monkeypatch -): +async def test_provider_credential_access_patch_bypass_forbidden(_connected_db, _patch_team_admin_lookup, monkeypatch): """Cursor BugBot regression: a team-admin can't sneak `access.teams` onto a PROVIDER credential to route through the decider instead of the admin gate. @@ -467,9 +503,7 @@ async def test_provider_credential_access_patch_bypass_forbidden( @pytest.mark.asyncio -async def test_team_admin_can_revoke_own_team_grant( - _connected_db, _patch_team_admin_lookup, monkeypatch -): +async def test_team_admin_can_revoke_own_team_grant(_connected_db, _patch_team_admin_lookup, monkeypatch): """A team-admin saving an access list without their own team_id revokes it.""" existing = CredentialItem( credential_name="dest", @@ -502,9 +536,7 @@ async def test_team_admin_can_revoke_own_team_grant( @pytest.mark.asyncio -async def test_team_admin_cannot_grant_foreign_team( - _connected_db, _patch_team_admin_lookup, monkeypatch -): +async def test_team_admin_cannot_grant_foreign_team(_connected_db, _patch_team_admin_lookup, monkeypatch): monkeypatch.setattr(litellm, "credential_list", [_resident_logging_dest()]) _patch_team_admin_lookup["ids"] = frozenset({"team-T"}) @@ -515,9 +547,7 @@ async def test_team_admin_cannot_grant_foreign_team( credential=CredentialItem( credential_name="dest", credential_values={}, - credential_info={ - "access": {"teams": ["team-existing", "team-foreign"]} - }, + credential_info={"access": {"teams": ["team-existing", "team-foreign"]}}, ), credential_name="dest", user_api_key_dict=_team_admin_of(["team-T"]), @@ -527,9 +557,7 @@ async def test_team_admin_cannot_grant_foreign_team( @pytest.mark.asyncio -async def test_team_admin_cannot_rotate_credential_values( - _connected_db, _patch_team_admin_lookup, monkeypatch -): +async def test_team_admin_cannot_rotate_credential_values(_connected_db, _patch_team_admin_lookup, monkeypatch): monkeypatch.setattr(litellm, "credential_list", [_resident_logging_dest()]) _patch_team_admin_lookup["ids"] = frozenset({"team-T"}) @@ -552,9 +580,7 @@ async def test_team_admin_cannot_rotate_credential_values( @pytest.mark.asyncio -async def test_team_admin_cannot_flip_global( - _connected_db, _patch_team_admin_lookup, monkeypatch -): +async def test_team_admin_cannot_flip_global(_connected_db, _patch_team_admin_lookup, monkeypatch): monkeypatch.setattr(litellm, "credential_list", [_resident_logging_dest()]) _patch_team_admin_lookup["ids"] = frozenset({"team-T"}) @@ -575,9 +601,7 @@ async def test_team_admin_cannot_flip_global( @pytest.mark.asyncio -async def test_get_credentials_shows_only_in_scope_destinations_for_non_admin( - monkeypatch, _patch_team_admin_lookup -): +async def test_get_credentials_shows_only_in_scope_destinations_for_non_admin(monkeypatch, _patch_team_admin_lookup): """Leak regression (Veria #1): a non-proxy-admin sees only destinations granted to a scope they administer, not every logging destination. The caller admins team-existing, so they see the destination granted to team-existing but never @@ -618,9 +642,7 @@ async def test_get_credentials_shows_only_in_scope_destinations_for_non_admin( @pytest.mark.asyncio -async def test_get_credentials_masks_otel_headers_only_for_non_admin( - monkeypatch, _patch_team_admin_lookup -): +async def test_get_credentials_masks_otel_headers_only_for_non_admin(monkeypatch, _patch_team_admin_lookup): raw_headers = "Authorization=Bearer collector-secret,x-api-key=api-secret" monkeypatch.setattr( litellm, @@ -656,9 +678,7 @@ async def test_get_credentials_masks_otel_headers_only_for_non_admin( @pytest.mark.asyncio -async def test_get_credentials_hides_out_of_scope_destination( - monkeypatch, _patch_team_admin_lookup -): +async def test_get_credentials_hides_out_of_scope_destination(monkeypatch, _patch_team_admin_lookup): """The exact leak: a team-admin of an unrelated team must see none of another team's destinations. Pre-fix, get_credentials returned every logging destination to any team/org admin regardless of the destination's access.""" @@ -683,9 +703,7 @@ async def test_get_credentials_hides_out_of_scope_destination( @pytest.mark.asyncio -async def test_get_credentials_shows_org_scoped_destination_to_org_admin( - monkeypatch, _patch_team_admin_lookup -): +async def test_get_credentials_shows_org_scoped_destination_to_org_admin(monkeypatch, _patch_team_admin_lookup): """An org-admin sees a destination granted to their org via access.orgs, matched against the org-admin scope (not just team ids).""" monkeypatch.setattr( @@ -708,18 +726,14 @@ async def test_get_credentials_shows_org_scoped_destination_to_org_admin( response = await endpoints.get_credentials( request=MagicMock(), fastapi_response=MagicMock(), - user_api_key_dict=UserAPIKeyAuth( - api_key="k", user_role=LitellmUserRoles.INTERNAL_USER, user_id="oa" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="k", user_role=LitellmUserRoles.INTERNAL_USER, user_id="oa"), ) names = [c["credential_name"] for c in response["credentials"]] assert names == ["org-dest"] @pytest.mark.asyncio -async def test_get_credentials_forbidden_for_plain_user( - monkeypatch, _patch_team_admin_lookup -): +async def test_get_credentials_forbidden_for_plain_user(monkeypatch, _patch_team_admin_lookup): """Veria F2 regression: a plain internal_user (no team-admin or org-admin status anywhere) gets 403, NOT a filtered list. The previous handler returned destination names, hosts, and scope metadata to any @@ -773,9 +787,7 @@ def test_patch_credentials_route_targets_update_credential(): @pytest.mark.asyncio -async def test_patch_credentials_does_not_leak_credential_type( - _connected_db, _patch_team_admin_lookup, monkeypatch -): +async def test_patch_credentials_does_not_leak_credential_type(_connected_db, _patch_team_admin_lookup, monkeypatch): """Existence-oracle regression: a team-admin probing a credential they don't own must NOT be able to distinguish "logging credential, not yours" from "provider credential" or "doesn't exist" by comparing 403 detail strings. @@ -847,9 +859,7 @@ async def test_patch_credentials_echoes_foreign_team_id_to_legit_team_admin( request=MagicMock(), fastapi_response=MagicMock(), credential=UpdateCredentialItem( - credential_info={ - "access": {"teams": ["team-existing", "team-foreign"]} - }, + credential_info={"access": {"teams": ["team-existing", "team-foreign"]}}, ), credential_name="dest", user_api_key_dict=_team_admin_of(["team-T"]), @@ -892,9 +902,7 @@ async def test_get_credentials_returns_all_for_proxy_admin(monkeypatch): ) names = sorted(c["credential_name"] for c in response["credentials"]) assert names == ["generic-otel", "openai", "poc-langfuse"] - generic = next( - c for c in response["credentials"] if c["credential_name"] == "generic-otel" - ) + generic = next(c for c in response["credentials"] if c["credential_name"] == "generic-otel") assert generic["credential_values"]["otel_headers"] == raw_headers