scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(proxy): harden credential scope and migration recovery

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-07-21 06:11:41 +00:00
parent b55f45bd38
commit 7b78e1426b
4 changed files with 196 additions and 264 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

190
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -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:

73
litellm-proxy-extras/tests/test_setup_database_fail_fast.py
View file

@ -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]

51
litellm/proxy/credential_endpoints/endpoints.py
View file

@ -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,

146
tests/test_litellm/proxy/credential_endpoints/test_endpoints.py
View file

@ -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