mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-19 22:19:07 +00:00
Merge pull request #25796 from BerriAI/litellm_yj_apr14
[Infra] Merge dev branch
This commit is contained in:
@@ -804,6 +804,8 @@ router_settings:
|
||||
| LITELLM_ASSETS_PATH | Path to directory for UI assets and logos. Used when running with read-only filesystem (e.g., Kubernetes). Default is `/var/lib/litellm/assets` in Docker.
|
||||
| LITELLM_BLOG_POSTS_URL | Custom URL for fetching LiteLLM blog posts JSON. Default is the GitHub main branch URL
|
||||
| LITELLM_CLI_JWT_EXPIRATION_HOURS | Expiration time in hours for CLI-generated JWT tokens. Default is 24 hours
|
||||
| LITELLM_CORS_ALLOW_CREDENTIALS | Set to `true` to explicitly allow credentials in CORS responses. When not set, credentials are disabled automatically if `LITELLM_CORS_ORIGINS` is `*` (wildcard) to prevent the browser security misconfiguration of reflecting any origin with credentials
|
||||
| LITELLM_CORS_ORIGINS | Comma-separated list of allowed CORS origins (e.g. `https://app.example.com,https://admin.example.com`). Defaults to `*` (all origins) when not set
|
||||
| LITELLM_DD_AGENT_HOST | Hostname or IP of DataDog agent for LiteLLM-specific logging. When set, logs are sent to agent instead of direct API
|
||||
| LITELLM_DEPLOYMENT_ENVIRONMENT | Environment name for the deployment (e.g., "production", "staging"). Used as a fallback when OTEL_ENVIRONMENT_NAME is not set. Sets the `environment` tag in telemetry data
|
||||
| LITELLM_DETAILED_TIMING | When true, adds detailed per-phase timing headers to responses (`x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms`). Default is false. See [latency overhead docs](../troubleshoot/latency_overhead.md)
|
||||
|
||||
@@ -4,6 +4,12 @@ from litellm import verbose_logger
|
||||
|
||||
_db = Any
|
||||
|
||||
# Markers that indicate a view/relation does not yet exist in the database.
|
||||
# Keeping these in one place avoids repeating the check across all view blocks
|
||||
# and prevents overly broad matches (e.g. bare 'undefined' would also match
|
||||
# 'undefined function' or 'column undefined_col referenced in query').
|
||||
_VIEW_NOT_FOUND_MARKERS = ("does not exist", "no such table", "undefined table")
|
||||
|
||||
|
||||
async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
"""
|
||||
@@ -18,14 +24,17 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
|
||||
If the view doesn't exist, one will be created.
|
||||
"""
|
||||
|
||||
try:
|
||||
# Try to select one row from the view
|
||||
await db.query_raw("""SELECT 1 FROM "LiteLLM_VerificationTokenView" LIMIT 1""")
|
||||
print("LiteLLM_VerificationTokenView Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("LiteLLM_VerificationTokenView Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
|
||||
raise
|
||||
# If an error occurs, the view does not exist, so create it
|
||||
await db.execute_raw(
|
||||
"""
|
||||
await db.execute_raw("""
|
||||
CREATE VIEW "LiteLLM_VerificationTokenView" AS
|
||||
SELECT
|
||||
v.*,
|
||||
@@ -37,15 +46,17 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
FROM "LiteLLM_VerificationToken" v
|
||||
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id
|
||||
LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id = p.project_id;
|
||||
"""
|
||||
)
|
||||
""")
|
||||
|
||||
print("LiteLLM_VerificationTokenView Created!") # noqa
|
||||
verbose_logger.debug("LiteLLM_VerificationTokenView Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpend" LIMIT 1""")
|
||||
print("MonthlyGlobalSpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("MonthlyGlobalSpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "MonthlyGlobalSpend" AS
|
||||
SELECT
|
||||
@@ -60,12 +71,15 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("MonthlyGlobalSpend Created!") # noqa
|
||||
verbose_logger.debug("MonthlyGlobalSpend Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "Last30dKeysBySpend" LIMIT 1""")
|
||||
print("Last30dKeysBySpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("Last30dKeysBySpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "Last30dKeysBySpend" AS
|
||||
SELECT
|
||||
@@ -88,12 +102,15 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("Last30dKeysBySpend Created!") # noqa
|
||||
verbose_logger.debug("Last30dKeysBySpend Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "Last30dModelsBySpend" LIMIT 1""")
|
||||
print("Last30dModelsBySpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("Last30dModelsBySpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "Last30dModelsBySpend" AS
|
||||
SELECT
|
||||
@@ -111,11 +128,14 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("Last30dModelsBySpend Created!") # noqa
|
||||
verbose_logger.debug("Last30dModelsBySpend Created!")
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerKey" LIMIT 1""")
|
||||
print("MonthlyGlobalSpendPerKey Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("MonthlyGlobalSpendPerKey Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "MonthlyGlobalSpendPerKey" AS
|
||||
SELECT
|
||||
@@ -132,13 +152,16 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("MonthlyGlobalSpendPerKey Created!") # noqa
|
||||
verbose_logger.debug("MonthlyGlobalSpendPerKey Created!")
|
||||
try:
|
||||
await db.query_raw(
|
||||
"""SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1"""
|
||||
)
|
||||
print("MonthlyGlobalSpendPerUserPerKey Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "MonthlyGlobalSpendPerUserPerKey" AS
|
||||
SELECT
|
||||
@@ -157,12 +180,15 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("MonthlyGlobalSpendPerUserPerKey Created!") # noqa
|
||||
verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "DailyTagSpend" LIMIT 1""")
|
||||
print("DailyTagSpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("DailyTagSpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "DailyTagSpend" AS
|
||||
SELECT
|
||||
@@ -175,12 +201,15 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("DailyTagSpend Created!") # noqa
|
||||
verbose_logger.debug("DailyTagSpend Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "Last30dTopEndUsersSpend" LIMIT 1""")
|
||||
print("Last30dTopEndUsersSpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("Last30dTopEndUsersSpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE VIEW "Last30dTopEndUsersSpend" AS
|
||||
SELECT end_user, COUNT(*) AS total_events, SUM(spend) AS total_spend
|
||||
@@ -193,7 +222,7 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
||||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("Last30dTopEndUsersSpend Created!") # noqa
|
||||
verbose_logger.debug("Last30dTopEndUsersSpend Created!")
|
||||
|
||||
return
|
||||
|
||||
|
||||
@@ -82,7 +82,7 @@ class SpendLogCleanup:
|
||||
break
|
||||
# Step 1: Find logs and delete them in one go without fetching to application
|
||||
# Delete in batches, limited by self.batch_size
|
||||
deleted_count = await prisma_client.db.execute_raw(
|
||||
deleted_result = await prisma_client.db.execute_raw(
|
||||
"""
|
||||
DELETE FROM "LiteLLM_SpendLogs"
|
||||
WHERE "request_id" IN (
|
||||
@@ -94,6 +94,17 @@ class SpendLogCleanup:
|
||||
cutoff_date,
|
||||
self.batch_size,
|
||||
)
|
||||
|
||||
deleted_count = 0
|
||||
if isinstance(deleted_result, int):
|
||||
deleted_count = deleted_result
|
||||
else:
|
||||
verbose_proxy_logger.error(
|
||||
f"Unexpected execute_raw return type for spend log cleanup: {type(deleted_result)}; "
|
||||
"aborting cleanup to avoid infinite loop"
|
||||
)
|
||||
break
|
||||
|
||||
verbose_proxy_logger.info(f"Deleted {deleted_count} logs in this batch")
|
||||
|
||||
if deleted_count == 0:
|
||||
|
||||
@@ -98,11 +98,20 @@ class TeamMemberPermissionChecks:
|
||||
)
|
||||
|
||||
# 5. Check if the team member has permissions for the endpoint
|
||||
TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
|
||||
team_member_object=key_assigned_user_in_team,
|
||||
team_table=team_table,
|
||||
route=route,
|
||||
has_permission = (
|
||||
TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
|
||||
team_member_object=key_assigned_user_in_team,
|
||||
team_table=team_table,
|
||||
route=route,
|
||||
)
|
||||
)
|
||||
if not has_permission:
|
||||
raise ProxyException(
|
||||
message=f"User {user_api_key_dict.user_id} does not belong to team {team_table.team_id}. Team-scoped key management endpoints can only be used for keys in your own team.",
|
||||
type=ProxyErrorTypes.team_member_permission_error,
|
||||
param=route,
|
||||
code=401,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def does_team_member_have_permissions_for_endpoint(
|
||||
|
||||
@@ -54,7 +54,7 @@ from litellm.constants import (
|
||||
LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
|
||||
LITELLM_UI_ALLOW_HEADERS,
|
||||
LITELLM_UI_SESSION_DURATION,
|
||||
DAILY_TAG_SPEND_BATCH_MULTIPLIER
|
||||
DAILY_TAG_SPEND_BATCH_MULTIPLIER,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_init_custom_logger_compatible_class,
|
||||
@@ -1139,7 +1139,54 @@ async def openai_exception_handler(request: Request, exc: ProxyException):
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
origins = ["*"]
|
||||
|
||||
|
||||
def _get_cors_config(
|
||||
cors_origins_env: Optional[str] = None,
|
||||
cors_credentials_env: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Compute CORS allowed origins and credentials flag from environment variables.
|
||||
|
||||
Extracted into a function so it can be unit-tested without reloading the module.
|
||||
|
||||
Args:
|
||||
cors_origins_env: Value of LITELLM_CORS_ORIGINS (defaults to os.getenv).
|
||||
cors_credentials_env: Value of LITELLM_CORS_ALLOW_CREDENTIALS (defaults to os.getenv).
|
||||
|
||||
Returns:
|
||||
Tuple[List[str], bool]: (origins, allow_credentials)
|
||||
"""
|
||||
_origins_raw = (
|
||||
cors_origins_env
|
||||
if cors_origins_env is not None
|
||||
else os.getenv("LITELLM_CORS_ORIGINS")
|
||||
)
|
||||
if _origins_raw is None or _origins_raw.strip() == "":
|
||||
computed_origins = ["*"]
|
||||
else:
|
||||
computed_origins = [o.strip() for o in _origins_raw.split(",") if o.strip()]
|
||||
|
||||
# Disable credentials by default when wildcard origins are used — combining
|
||||
# allow_origins=["*"] with allow_credentials=True causes Starlette to reflect
|
||||
# the incoming Origin header, allowing any site to make credentialed requests.
|
||||
# Set LITELLM_CORS_ALLOW_CREDENTIALS=true to explicitly restore the old behaviour
|
||||
# (e.g. for non-browser clients that relied on the Access-Control-Allow-Credentials
|
||||
# header being present regardless of origin).
|
||||
_credentials_raw = (
|
||||
cors_credentials_env
|
||||
if cors_credentials_env is not None
|
||||
else os.getenv("LITELLM_CORS_ALLOW_CREDENTIALS")
|
||||
)
|
||||
if _credentials_raw is not None:
|
||||
computed_credentials = _credentials_raw.strip().lower() == "true"
|
||||
else:
|
||||
computed_credentials = "*" not in computed_origins
|
||||
|
||||
return computed_origins, computed_credentials
|
||||
|
||||
|
||||
origins, allow_cors_credentials = _get_cors_config()
|
||||
|
||||
|
||||
# get current directory
|
||||
@@ -1466,7 +1513,7 @@ current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=origins,
|
||||
allow_credentials=True,
|
||||
allow_credentials=allow_cors_credentials,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
expose_headers=LITELLM_UI_ALLOW_HEADERS,
|
||||
@@ -2315,9 +2362,13 @@ def _write_health_state_to_router_cache(
|
||||
|
||||
exception_status = getattr(original_exception, "status_code", 500)
|
||||
|
||||
if llm_router.health_check_ignore_transient_errors and exception_status in (
|
||||
429,
|
||||
408,
|
||||
if (
|
||||
llm_router.health_check_ignore_transient_errors
|
||||
and exception_status
|
||||
in (
|
||||
429,
|
||||
408,
|
||||
)
|
||||
):
|
||||
continue
|
||||
|
||||
@@ -6286,7 +6337,9 @@ class ProxyStartupEvent:
|
||||
|
||||
### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ###
|
||||
## Reduces QPS as there are more tags for a single request
|
||||
tag_spend_update_interval = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER)
|
||||
tag_spend_update_interval = int(
|
||||
batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER
|
||||
)
|
||||
from litellm.proxy.utils import update_daily_tag_spend
|
||||
|
||||
scheduler.add_job(
|
||||
@@ -7131,9 +7184,9 @@ async def chat_completion( # noqa: PLR0915
|
||||
hasattr(user_api_key_dict, "organization_alias")
|
||||
and user_api_key_dict.organization_alias is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_alias"] = (
|
||||
user_api_key_dict.organization_alias
|
||||
)
|
||||
data["metadata"][
|
||||
"user_api_key_org_alias"
|
||||
] = user_api_key_dict.organization_alias
|
||||
if (
|
||||
hasattr(user_api_key_dict, "agent_id")
|
||||
and user_api_key_dict.agent_id is not None
|
||||
@@ -7312,9 +7365,9 @@ async def completion( # noqa: PLR0915
|
||||
hasattr(user_api_key_dict, "organization_alias")
|
||||
and user_api_key_dict.organization_alias is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_alias"] = (
|
||||
user_api_key_dict.organization_alias
|
||||
)
|
||||
data["metadata"][
|
||||
"user_api_key_org_alias"
|
||||
] = user_api_key_dict.organization_alias
|
||||
if (
|
||||
hasattr(user_api_key_dict, "agent_id")
|
||||
and user_api_key_dict.agent_id is not None
|
||||
@@ -7561,9 +7614,9 @@ async def embeddings( # noqa: PLR0915
|
||||
hasattr(user_api_key_dict, "organization_alias")
|
||||
and user_api_key_dict.organization_alias is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_alias"] = (
|
||||
user_api_key_dict.organization_alias
|
||||
)
|
||||
data["metadata"][
|
||||
"user_api_key_org_alias"
|
||||
] = user_api_key_dict.organization_alias
|
||||
if (
|
||||
hasattr(user_api_key_dict, "agent_id")
|
||||
and user_api_key_dict.agent_id is not None
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
Tests for create_missing_views exception handling fix.
|
||||
|
||||
Verifies that real DB errors (auth failures, connection errors, etc.)
|
||||
are re-raised instead of being silently swallowed, while genuine
|
||||
"view not found" errors still trigger view creation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, call
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_reraises_connection_error():
|
||||
"""should re-raise exceptions that are NOT 'does not exist' errors (e.g. connection errors)."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=Exception("connection refused: unable to connect to database")
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock()
|
||||
|
||||
with pytest.raises(Exception, match="connection refused"):
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_reraises_permission_error():
|
||||
"""should re-raise permission denied errors, not treat them as missing views."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=Exception(
|
||||
"permission denied for table LiteLLM_VerificationTokenView"
|
||||
)
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock()
|
||||
|
||||
with pytest.raises(Exception, match="permission denied"):
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_creates_view_on_does_not_exist():
|
||||
"""should call execute_raw to create view when error contains 'does not exist'."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=[
|
||||
Exception('relation "LiteLLM_VerificationTokenView" does not exist'),
|
||||
None, # MonthlyGlobalSpend exists
|
||||
None, # Last30dKeysBySpend exists
|
||||
None, # Last30dModelsBySpend exists
|
||||
None, # MonthlyGlobalSpendPerKey exists
|
||||
None, # MonthlyGlobalSpendPerUserPerKey exists
|
||||
None, # DailyTagSpend exists
|
||||
None, # Last30dTopEndUsersSpend exists
|
||||
]
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock(return_value=None)
|
||||
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_called_once()
|
||||
created_sql = mock_db.execute_raw.call_args[0][0]
|
||||
assert 'CREATE VIEW "LiteLLM_VerificationTokenView"' in created_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_creates_view_on_undefined_error():
|
||||
"""should treat 'undefined' errors as 'view not found' and attempt creation."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=[
|
||||
Exception("undefined table LiteLLM_VerificationTokenView"),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
]
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock(return_value=None)
|
||||
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_skips_creation_when_view_exists():
|
||||
"""should not call execute_raw when all views already exist."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(return_value=[{"?column?": 1}])
|
||||
mock_db.execute_raw = AsyncMock()
|
||||
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_reraises_undefined_function_error():
|
||||
"""should re-raise 'undefined function' errors — bare 'undefined' is too broad
|
||||
and would previously misclassify DB function errors as missing-view signals."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=Exception("ERROR: undefined function pg_get_viewdef()")
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock()
|
||||
|
||||
with pytest.raises(Exception, match="undefined function"):
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_creates_view_on_undefined_table_error():
|
||||
"""should treat 'undefined table' as a missing-view signal and attempt creation."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=[
|
||||
Exception('undefined table "LiteLLM_VerificationTokenView"'),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
]
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock(return_value=None)
|
||||
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_called_once()
|
||||
@@ -8,7 +8,7 @@ sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.proxy._types import KeyManagementRoutes, Member
|
||||
from litellm.proxy._types import KeyManagementRoutes, Member, ProxyException
|
||||
from litellm.proxy.management_helpers.team_member_permission_checks import (
|
||||
BASELINE_TEAM_MEMBER_PERMISSIONS,
|
||||
TeamMemberPermissionChecks,
|
||||
@@ -188,3 +188,74 @@ class TestGetDefaultTeamParam:
|
||||
assert _get_default_team_param("budget_duration") == "7d"
|
||||
assert _get_default_team_param("tpm_limit") == 1000
|
||||
assert _get_default_team_param("rpm_limit") == 100
|
||||
|
||||
|
||||
class TestCanTeamMemberExecuteKeyManagementEndpoint:
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_user_not_in_keys_team(self, monkeypatch):
|
||||
"""Non-members should be blocked from team-scoped key management endpoints."""
|
||||
from litellm.proxy.management_endpoints import key_management_endpoints
|
||||
from litellm.proxy.management_helpers import team_member_permission_checks as module
|
||||
|
||||
async def _mock_get_team_object(**kwargs):
|
||||
team = MagicMock()
|
||||
team.team_id = "team-b"
|
||||
team.team_member_permissions = ["/key/update"]
|
||||
return team
|
||||
|
||||
monkeypatch.setattr(module, "get_team_object", _mock_get_team_object)
|
||||
monkeypatch.setattr(key_management_endpoints, "_get_user_in_team", lambda **kwargs: None)
|
||||
|
||||
user_api_key_dict = MagicMock()
|
||||
user_api_key_dict.user_role = "internal_user"
|
||||
user_api_key_dict.user_id = "user-a"
|
||||
user_api_key_dict.parent_otel_span = None
|
||||
|
||||
existing_key_row = MagicMock()
|
||||
existing_key_row.team_id = "team-b"
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route=KeyManagementRoutes.KEY_UPDATE,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
existing_key_row=existing_key_row,
|
||||
)
|
||||
assert str(exc.value.code) == "401"
|
||||
assert exc.value.type == "team_member_permission_error"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_allows_team_admin_in_keys_team(self, monkeypatch):
|
||||
"""Team admins of the key's team should be allowed."""
|
||||
from litellm.proxy.management_endpoints import key_management_endpoints
|
||||
from litellm.proxy.management_helpers import team_member_permission_checks as module
|
||||
|
||||
async def _mock_get_team_object(**kwargs):
|
||||
team = MagicMock()
|
||||
team.team_id = "team-a"
|
||||
team.team_member_permissions = ["/key/update"]
|
||||
return team
|
||||
|
||||
monkeypatch.setattr(module, "get_team_object", _mock_get_team_object)
|
||||
monkeypatch.setattr(
|
||||
key_management_endpoints,
|
||||
"_get_user_in_team",
|
||||
lambda **kwargs: Member(role="admin", user_id="user-a"),
|
||||
)
|
||||
|
||||
user_api_key_dict = MagicMock()
|
||||
user_api_key_dict.user_role = "internal_user"
|
||||
user_api_key_dict.user_id = "user-a"
|
||||
user_api_key_dict.parent_otel_span = None
|
||||
|
||||
existing_key_row = MagicMock()
|
||||
existing_key_row.team_id = "team-a"
|
||||
|
||||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route=KeyManagementRoutes.KEY_UPDATE,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
existing_key_row=existing_key_row,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
"""
|
||||
Tests for CORS configuration security fix.
|
||||
|
||||
All tests import _get_cors_config directly from proxy_server so they exercise
|
||||
real production code rather than a local mirror.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_cors_wildcard_disables_credentials():
|
||||
"""should disable credentials when LITELLM_CORS_ORIGINS is not set (defaults to wildcard)."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
origins, allow_credentials = _get_cors_config(cors_origins_env="")
|
||||
assert origins == ["*"]
|
||||
assert allow_credentials is False
|
||||
|
||||
|
||||
def test_cors_empty_string_disables_credentials():
|
||||
"""should disable credentials when LITELLM_CORS_ORIGINS is empty or whitespace."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
for empty in ("", " ", "\t"):
|
||||
origins, allow_credentials = _get_cors_config(cors_origins_env=empty)
|
||||
assert origins == ["*"], f"Expected wildcard for input {repr(empty)}"
|
||||
assert (
|
||||
allow_credentials is False
|
||||
), f"Expected no credentials for input {repr(empty)}"
|
||||
|
||||
|
||||
def test_cors_single_specific_origin_enables_credentials():
|
||||
"""should enable credentials when a single explicit origin is configured."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
origins, allow_credentials = _get_cors_config(
|
||||
cors_origins_env="https://admin.example.com"
|
||||
)
|
||||
assert origins == ["https://admin.example.com"]
|
||||
assert allow_credentials is True
|
||||
|
||||
|
||||
def test_cors_multiple_specific_origins_enables_credentials():
|
||||
"""should enable credentials and correctly parse comma-separated origins."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
origins, allow_credentials = _get_cors_config(
|
||||
cors_origins_env="https://app.example.com, https://admin.example.com, https://api.example.com"
|
||||
)
|
||||
assert origins == [
|
||||
"https://app.example.com",
|
||||
"https://admin.example.com",
|
||||
"https://api.example.com",
|
||||
]
|
||||
assert allow_credentials is True
|
||||
|
||||
|
||||
def test_cors_wildcard_string_in_env_disables_credentials():
|
||||
"""should disable credentials when LITELLM_CORS_ORIGINS is explicitly set to '*'."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
origins, allow_credentials = _get_cors_config(cors_origins_env="*")
|
||||
assert "*" in origins
|
||||
assert allow_credentials is False
|
||||
|
||||
|
||||
def test_cors_origins_strips_whitespace():
|
||||
"""should strip surrounding whitespace from each origin entry."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
origins, _ = _get_cors_config(
|
||||
cors_origins_env=" https://a.com , https://b.com "
|
||||
)
|
||||
assert origins == ["https://a.com", "https://b.com"]
|
||||
|
||||
|
||||
def test_cors_origins_skips_blank_entries():
|
||||
"""should skip blank entries caused by trailing/double commas."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
origins, allow_credentials = _get_cors_config(
|
||||
cors_origins_env="https://a.com,,https://b.com,"
|
||||
)
|
||||
assert origins == ["https://a.com", "https://b.com"]
|
||||
assert allow_credentials is True
|
||||
|
||||
|
||||
def test_cors_explicit_credentials_true_overrides_wildcard():
|
||||
"""should enable credentials when LITELLM_CORS_ALLOW_CREDENTIALS=true even
|
||||
if wildcard origins are in use (opt-in for existing deployments)."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
origins, allow_credentials = _get_cors_config(
|
||||
cors_origins_env="",
|
||||
cors_credentials_env="true",
|
||||
)
|
||||
assert "*" in origins
|
||||
assert allow_credentials is True
|
||||
|
||||
|
||||
def test_cors_explicit_credentials_false_overrides_specific_origins():
|
||||
"""should disable credentials when LITELLM_CORS_ALLOW_CREDENTIALS=false even
|
||||
if specific origins are configured."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
origins, allow_credentials = _get_cors_config(
|
||||
cors_origins_env="https://admin.example.com",
|
||||
cors_credentials_env="false",
|
||||
)
|
||||
assert origins == ["https://admin.example.com"]
|
||||
assert allow_credentials is False
|
||||
|
||||
|
||||
def test_cors_explicit_credentials_case_insensitive():
|
||||
"""should accept TRUE/FALSE case-insensitively for LITELLM_CORS_ALLOW_CREDENTIALS."""
|
||||
from litellm.proxy.proxy_server import _get_cors_config
|
||||
|
||||
_, allow_true = _get_cors_config(cors_origins_env="", cors_credentials_env="TRUE")
|
||||
_, allow_false = _get_cors_config(
|
||||
cors_origins_env="https://x.com", cors_credentials_env="FALSE"
|
||||
)
|
||||
assert allow_true is True
|
||||
assert allow_false is False
|
||||
|
||||
|
||||
def test_proxy_server_cors_invariant():
|
||||
"""should verify that proxy_server module-level origins and allow_cors_credentials
|
||||
are consistent — catches any future drift in the module-level call to _get_cors_config.
|
||||
"""
|
||||
import os
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
if os.getenv("LITELLM_CORS_ALLOW_CREDENTIALS") is None:
|
||||
assert proxy_server.allow_cors_credentials == (
|
||||
"*" not in proxy_server.origins
|
||||
), (
|
||||
f"Invariant broken: allow_cors_credentials={proxy_server.allow_cors_credentials} "
|
||||
f"but origins={proxy_server.origins}. "
|
||||
"When origins contains '*', allow_credentials must be False."
|
||||
)
|
||||
@@ -287,9 +287,48 @@ def test_string_retention_still_works():
|
||||
general_settings={"maximum_spend_logs_retention_period": setting}
|
||||
)
|
||||
assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}"
|
||||
assert cleaner.retention_seconds == expected_seconds, (
|
||||
f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}"
|
||||
)
|
||||
assert (
|
||||
cleaner.retention_seconds == expected_seconds
|
||||
), f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_old_logs_aborts_on_non_int_execute_raw_return():
|
||||
"""should abort deletion loop immediately when execute_raw returns a non-int
|
||||
(e.g. None or dict), preventing an infinite loop."""
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_db.execute_raw = AsyncMock(return_value=None)
|
||||
mock_prisma_client.db = mock_db
|
||||
|
||||
cleaner = SpendLogCleanup(
|
||||
general_settings={"maximum_spend_logs_retention_period": "7d"}
|
||||
)
|
||||
|
||||
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
||||
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
|
||||
|
||||
assert mock_db.execute_raw.call_count == 1
|
||||
assert total_deleted == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_old_logs_continues_on_valid_int_return():
|
||||
"""should continue deletion loop across batches when execute_raw returns valid int counts."""
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0])
|
||||
mock_prisma_client.db = mock_db
|
||||
|
||||
cleaner = SpendLogCleanup(
|
||||
general_settings={"maximum_spend_logs_retention_period": "7d"}
|
||||
)
|
||||
|
||||
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
||||
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
|
||||
|
||||
assert mock_db.execute_raw.call_count == 3
|
||||
assert total_deleted == 800
|
||||
|
||||
|
||||
def test_cleanup_batch_size_env_var(monkeypatch):
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
import React from "react";
|
||||
import { screen, waitFor, within } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { renderWithProviders } from "../../../tests/test-utils";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import PoliciesPanel from "./index";
|
||||
|
||||
/**
|
||||
* Ant Design's static Modal.confirm often does not run onOk in the real app (React 18+).
|
||||
* In jsdom it may still run; we mock confirm as a no-op so the test fails until the panel
|
||||
* uses a controlled DeleteResourceModal instead of Modal.confirm.
|
||||
*/
|
||||
vi.mock("antd", async (importOriginal) => {
|
||||
const mod = await importOriginal<typeof import("antd")>();
|
||||
return {
|
||||
...mod,
|
||||
Modal: Object.assign(mod.Modal, {
|
||||
confirm: vi.fn(),
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
const EXPECTED_ATTACHMENT_ID = "att-11111111-2222-3333-4444-555555555555" as const;
|
||||
|
||||
const networkingMocks = vi.hoisted(() => ({
|
||||
deletePolicyAttachmentCall: vi.fn().mockResolvedValue(undefined),
|
||||
getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }),
|
||||
getPolicyAttachmentsList: vi.fn().mockResolvedValue({
|
||||
attachments: [
|
||||
{
|
||||
attachment_id: "att-11111111-2222-3333-4444-555555555555",
|
||||
policy_name: "test-policy",
|
||||
scope: null,
|
||||
teams: [],
|
||||
keys: [],
|
||||
models: [],
|
||||
tags: [],
|
||||
},
|
||||
],
|
||||
}),
|
||||
getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }),
|
||||
getPolicyInfo: vi.fn().mockResolvedValue({}),
|
||||
deletePolicyCall: vi.fn().mockResolvedValue(undefined),
|
||||
createPolicyCall: vi.fn(),
|
||||
updatePolicyCall: vi.fn(),
|
||||
createPolicyAttachmentCall: vi.fn(),
|
||||
createGuardrailCall: vi.fn(),
|
||||
enrichPolicyTemplate: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("../networking", () => ({
|
||||
...networkingMocks,
|
||||
}));
|
||||
|
||||
vi.mock("./impact_popover", () => ({
|
||||
default: () => <button type="button" aria-label="View blast radius" />,
|
||||
}));
|
||||
|
||||
vi.mock("@heroicons/react/outline", () => ({
|
||||
TrashIcon: function TrashIcon() {
|
||||
return null;
|
||||
},
|
||||
SwitchVerticalIcon: function SwitchVerticalIcon() {
|
||||
return null;
|
||||
},
|
||||
ChevronUpIcon: function ChevronUpIcon() {
|
||||
return null;
|
||||
},
|
||||
ChevronDownIcon: function ChevronDownIcon() {
|
||||
return null;
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock("@tremor/react", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("@tremor/react")>();
|
||||
return {
|
||||
...actual,
|
||||
Button: React.forwardRef<HTMLButtonElement, any>(({ children, ...props }, ref) =>
|
||||
React.createElement("button", { ...props, ref }, children),
|
||||
),
|
||||
Tooltip: ({ children }: { children?: React.ReactNode }) =>
|
||||
React.createElement(React.Fragment, null, children),
|
||||
Switch: ({
|
||||
checked,
|
||||
onChange,
|
||||
className,
|
||||
}: {
|
||||
checked?: boolean;
|
||||
onChange?: (v: boolean) => void;
|
||||
className?: string;
|
||||
}) =>
|
||||
React.createElement("input", {
|
||||
type: "checkbox",
|
||||
role: "switch",
|
||||
checked,
|
||||
onChange: (e: React.ChangeEvent<HTMLInputElement>) => onChange?.(e.target.checked),
|
||||
className,
|
||||
}),
|
||||
Icon: ({ icon: _IconComp, onClick, className }: any) =>
|
||||
React.createElement(
|
||||
"button",
|
||||
{ type: "button", onClick, className },
|
||||
"TrashIcon",
|
||||
),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock("./policy_templates", () => ({
|
||||
__esModule: true,
|
||||
default: () => <div data-testid="policy-templates-stub" />,
|
||||
}));
|
||||
|
||||
vi.mock("./pipeline_flow_builder", () => ({
|
||||
FlowBuilderPage: () => null,
|
||||
}));
|
||||
|
||||
vi.mock("./policy_info", () => ({
|
||||
__esModule: true,
|
||||
default: () => null,
|
||||
}));
|
||||
|
||||
vi.mock("./add_policy_form", () => ({
|
||||
__esModule: true,
|
||||
default: () => null,
|
||||
}));
|
||||
|
||||
vi.mock("./guardrail_selection_modal", () => ({
|
||||
__esModule: true,
|
||||
default: () => null,
|
||||
}));
|
||||
|
||||
vi.mock("./template_parameter_modal", () => ({
|
||||
__esModule: true,
|
||||
default: () => null,
|
||||
}));
|
||||
|
||||
vi.mock("./ai_suggestion_modal", () => ({
|
||||
__esModule: true,
|
||||
default: () => null,
|
||||
}));
|
||||
|
||||
vi.mock("./policy_test_panel", () => ({
|
||||
__esModule: true,
|
||||
default: () => null,
|
||||
}));
|
||||
|
||||
vi.mock("./add_attachment_form", () => ({
|
||||
__esModule: true,
|
||||
default: () => null,
|
||||
}));
|
||||
|
||||
describe("PoliciesPanel attachment delete", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("should call deletePolicyAttachmentCall after the user confirms delete in the attachment modal", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<PoliciesPanel accessToken="test-token" userRole="Admin" />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networkingMocks.getPolicyAttachmentsList).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
await user.click(screen.getByRole("tab", { name: /^attachments$/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("test-policy")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /TrashIcon/i }));
|
||||
|
||||
const dialog = await screen.findByRole("dialog", {}, { timeout: 5000 });
|
||||
expect(
|
||||
within(dialog).getByText(/Are you sure you want to delete this attachment/i),
|
||||
).toBeInTheDocument();
|
||||
|
||||
await user.click(within(dialog).getByRole("button", { name: /^delete$/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networkingMocks.deletePolicyAttachmentCall).toHaveBeenCalledTimes(1);
|
||||
expect(networkingMocks.deletePolicyAttachmentCall).toHaveBeenCalledWith("test-token", EXPECTED_ATTACHMENT_ID);
|
||||
});
|
||||
});
|
||||
|
||||
it("should show mutation pending state while attachment delete is in flight", async () => {
|
||||
let resolveDelete: (() => void) | undefined;
|
||||
const deletePromise = new Promise<void>((resolve) => {
|
||||
resolveDelete = resolve;
|
||||
});
|
||||
networkingMocks.deletePolicyAttachmentCall.mockImplementationOnce(() => deletePromise);
|
||||
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<PoliciesPanel accessToken="test-token" userRole="Admin" />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networkingMocks.getPolicyAttachmentsList).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
await user.click(screen.getByRole("tab", { name: /^attachments$/i }));
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("test-policy")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /TrashIcon/i }));
|
||||
const dialog = await screen.findByRole("dialog", {}, { timeout: 5000 });
|
||||
|
||||
const deleteButton = within(dialog).getByRole("button", { name: /^delete$/i });
|
||||
await user.click(deleteButton);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(within(dialog).getByRole("button", { name: /deleting/i })).toBeDisabled();
|
||||
});
|
||||
|
||||
resolveDelete?.();
|
||||
await waitFor(() => {
|
||||
expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,8 +1,9 @@
|
||||
import React, { useState, useEffect, useCallback } from "react";
|
||||
import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react";
|
||||
import { Modal, Alert } from "antd";
|
||||
import { Alert } from "antd";
|
||||
|
||||
import MessageManager from "@/components/molecules/message_manager";
|
||||
import { ExclamationCircleOutlined, InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { isAdminRole } from "@/utils/roles";
|
||||
import PolicyTable from "./policy_table";
|
||||
import PolicyInfoView from "./policy_info";
|
||||
@@ -15,11 +16,11 @@ import PolicyTemplates from "./policy_templates";
|
||||
import GuardrailSelectionModal from "./guardrail_selection_modal";
|
||||
import TemplateParameterModal from "./template_parameter_modal";
|
||||
import AiSuggestionModal from "./ai_suggestion_modal";
|
||||
import { useDeletePolicyAttachment } from "@/hooks/policies/useDeletePolicyAttachment";
|
||||
import {
|
||||
getPoliciesList,
|
||||
deletePolicyCall,
|
||||
getPolicyAttachmentsList,
|
||||
deletePolicyAttachmentCall,
|
||||
getGuardrailsList,
|
||||
getPolicyInfo,
|
||||
createPolicyCall,
|
||||
@@ -57,6 +58,8 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
|
||||
const [isDeleting, setIsDeleting] = useState(false);
|
||||
const [policyToDelete, setPolicyToDelete] = useState<Policy | null>(null);
|
||||
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
|
||||
const [attachmentToDelete, setAttachmentToDelete] = useState<PolicyAttachment | null>(null);
|
||||
const [isDeleteAttachmentModalOpen, setIsDeleteAttachmentModalOpen] = useState(false);
|
||||
const [isGuardrailSelectionModalOpen, setIsGuardrailSelectionModalOpen] = useState(false);
|
||||
const [selectedTemplate, setSelectedTemplate] = useState<any>(null);
|
||||
const [existingGuardrailNames, setExistingGuardrailNames] = useState<Set<string>>(new Set());
|
||||
@@ -166,24 +169,28 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
|
||||
setPolicyToDelete(null);
|
||||
};
|
||||
|
||||
const handleDeleteAttachment = (attachmentId: string) => {
|
||||
Modal.confirm({
|
||||
title: "Delete Attachment",
|
||||
icon: <ExclamationCircleOutlined />,
|
||||
content: "Are you sure you want to delete this attachment? This action cannot be undone.",
|
||||
okText: "Delete",
|
||||
okType: "danger",
|
||||
cancelText: "Cancel",
|
||||
onOk: async () => {
|
||||
if (!accessToken) return;
|
||||
try {
|
||||
await deletePolicyAttachmentCall(accessToken, attachmentId);
|
||||
MessageManager.success("Attachment deleted successfully");
|
||||
fetchAttachments();
|
||||
} catch (error) {
|
||||
console.error("Error deleting attachment:", error);
|
||||
MessageManager.error("Failed to delete attachment");
|
||||
}
|
||||
const deleteAttachmentMutation = useDeletePolicyAttachment({
|
||||
accessToken,
|
||||
onSuccess: fetchAttachments,
|
||||
});
|
||||
|
||||
const handleDeleteAttachmentClick = (attachmentId: string) => {
|
||||
const attachment = attachmentsList.find((a) => a.attachment_id === attachmentId) || null;
|
||||
setAttachmentToDelete(attachment);
|
||||
setIsDeleteAttachmentModalOpen(true);
|
||||
};
|
||||
|
||||
const handleAttachmentDeleteCancel = () => {
|
||||
setIsDeleteAttachmentModalOpen(false);
|
||||
setAttachmentToDelete(null);
|
||||
};
|
||||
|
||||
const handleAttachmentDeleteConfirm = () => {
|
||||
if (!attachmentToDelete) return;
|
||||
deleteAttachmentMutation.mutate(attachmentToDelete.attachment_id, {
|
||||
onSettled: () => {
|
||||
setIsDeleteAttachmentModalOpen(false);
|
||||
setAttachmentToDelete(null);
|
||||
},
|
||||
});
|
||||
};
|
||||
@@ -579,7 +586,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
|
||||
<AttachmentTable
|
||||
attachments={attachmentsList}
|
||||
isLoading={isAttachmentsLoading}
|
||||
onDeleteClick={handleDeleteAttachment}
|
||||
onDeleteClick={handleDeleteAttachmentClick}
|
||||
isAdmin={isAdmin}
|
||||
accessToken={accessToken}
|
||||
/>
|
||||
@@ -600,6 +607,21 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
|
||||
</TabPanels>
|
||||
</TabGroup>
|
||||
|
||||
<DeleteResourceModal
|
||||
isOpen={isDeleteAttachmentModalOpen}
|
||||
title="Delete Attachment"
|
||||
message="Are you sure you want to delete this attachment? This action cannot be undone."
|
||||
resourceInformationTitle="Attachment Information"
|
||||
resourceInformation={[
|
||||
{ label: "Attachment ID", value: attachmentToDelete?.attachment_id, code: true },
|
||||
{ label: "Policy", value: attachmentToDelete?.policy_name ?? "-" },
|
||||
{ label: "Scope", value: attachmentToDelete?.scope ?? "-" },
|
||||
]}
|
||||
onCancel={handleAttachmentDeleteCancel}
|
||||
onOk={handleAttachmentDeleteConfirm}
|
||||
confirmLoading={deleteAttachmentMutation.isPending}
|
||||
/>
|
||||
|
||||
<AiSuggestionModal
|
||||
visible={isAiSuggestionModalOpen}
|
||||
onSelectTemplates={(selectedTemplates) => {
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
import React from "react";
|
||||
import { renderHook, waitFor } from "@testing-library/react";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { useDeletePolicyAttachment } from "./useDeletePolicyAttachment";
|
||||
import { deletePolicyAttachmentCall } from "@/components/networking";
|
||||
import MessageManager from "@/components/molecules/message_manager";
|
||||
import { vi, describe, beforeEach, it, expect } from "vitest";
|
||||
|
||||
// Mock dependencies
|
||||
vi.mock("@/components/networking", () => ({
|
||||
deletePolicyAttachmentCall: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/molecules/message_manager", () => ({
|
||||
default: {
|
||||
success: vi.fn(),
|
||||
error: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
describe("useDeletePolicyAttachment", () => {
|
||||
let queryClient: QueryClient;
|
||||
|
||||
beforeEach(() => {
|
||||
queryClient = new QueryClient();
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
const wrapper = ({ children }: { children: React.ReactNode }) => (
|
||||
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
|
||||
);
|
||||
|
||||
it("should successfully delete a policy attachment and call onSuccess", async () => {
|
||||
const mockOnSuccess = vi.fn();
|
||||
(deletePolicyAttachmentCall as any).mockResolvedValue({});
|
||||
|
||||
const { result } = renderHook(
|
||||
() =>
|
||||
useDeletePolicyAttachment({
|
||||
accessToken: "test-token",
|
||||
onSuccess: mockOnSuccess,
|
||||
}),
|
||||
{ wrapper }
|
||||
);
|
||||
|
||||
result.current.mutate("attachment-1");
|
||||
|
||||
await waitFor(() => {
|
||||
expect(result.current.isSuccess).toBe(true);
|
||||
});
|
||||
|
||||
expect(deletePolicyAttachmentCall).toHaveBeenCalledWith("test-token", "attachment-1");
|
||||
expect(MessageManager.success).toHaveBeenCalledWith("Attachment deleted successfully");
|
||||
expect(mockOnSuccess).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should handle error when deleting policy attachment", async () => {
|
||||
const mockOnError = vi.fn();
|
||||
const error = new Error("Delete failed");
|
||||
(deletePolicyAttachmentCall as any).mockRejectedValue(error);
|
||||
|
||||
const { result } = renderHook(
|
||||
() =>
|
||||
useDeletePolicyAttachment({
|
||||
accessToken: "test-token",
|
||||
onError: mockOnError,
|
||||
}),
|
||||
{ wrapper }
|
||||
);
|
||||
|
||||
result.current.mutate("attachment-1");
|
||||
|
||||
await waitFor(() => {
|
||||
expect(result.current.isError).toBe(true);
|
||||
});
|
||||
|
||||
expect(deletePolicyAttachmentCall).toHaveBeenCalledWith("test-token", "attachment-1");
|
||||
expect(MessageManager.error).toHaveBeenCalledWith("Failed to delete attachment");
|
||||
expect(mockOnError).toHaveBeenCalledWith(error);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,37 @@
|
||||
import { useMutation } from "@tanstack/react-query";
|
||||
import { deletePolicyAttachmentCall } from "@/components/networking";
|
||||
import MessageManager from "@/components/molecules/message_manager";
|
||||
|
||||
interface UseDeletePolicyAttachmentProps {
|
||||
accessToken: string | null;
|
||||
onSuccess?: () => void;
|
||||
onError?: (error: any) => void;
|
||||
}
|
||||
|
||||
export const useDeletePolicyAttachment = ({
|
||||
accessToken,
|
||||
onSuccess,
|
||||
onError,
|
||||
}: UseDeletePolicyAttachmentProps) => {
|
||||
return useMutation({
|
||||
mutationFn: async (attachmentId: string) => {
|
||||
if (!accessToken) {
|
||||
throw new Error("Access token is required");
|
||||
}
|
||||
return deletePolicyAttachmentCall(accessToken, attachmentId);
|
||||
},
|
||||
onSuccess: () => {
|
||||
MessageManager.success("Attachment deleted successfully");
|
||||
if (onSuccess) {
|
||||
onSuccess();
|
||||
}
|
||||
},
|
||||
onError: (error) => {
|
||||
console.error("Error deleting attachment:", error);
|
||||
MessageManager.error("Failed to delete attachment");
|
||||
if (onError) {
|
||||
onError(error);
|
||||
}
|
||||
},
|
||||
});
|
||||
};
|
||||
Reference in New Issue
Block a user