diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index c12a90ed6a2..fdf2ba11192 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -12,9 +12,16 @@ All /tag management endpoints import asyncio import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Final, Protocol, TypedDict, overload +from typing import ( + TYPE_CHECKING, + Final, + Protocol, + TypedDict, + cast, # noqa: TID251 # casts over dynamic prisma Json storage + overload, +) from fastapi import APIRouter, Depends, HTTPException, Query from pydantic import TypeAdapter @@ -31,6 +38,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity, ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity +from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ( DailyTagSpendRepository, @@ -321,19 +329,41 @@ async def new_tag( # Update models with new tag if tag.models: - tasks: Final = [] + tasks: Final[list[Awaitable[object]]] = [] # mutable-ok: local accumulator, awaited before return + task_deployments: Final[list[Deployment]] = [] # mutable-ok: parallel to tasks, read after gather for model in tag.models: deployments = await get_deployments_by_model(model, llm_router) - tasks.extend( - [ + for deployment in deployments: + task_deployments.append(deployment) + tasks.append( _add_tag_to_deployment( deployment=deployment, tag=tag.name, ) - for deployment in deployments - ] + ) + # return_exceptions so every deployment write settles before we look: + # with the default, one failure would leave sibling tasks running + # while we roll back underneath them. + tag_results: Final = await asyncio.gather(*tasks, return_exceptions=True) + tag_error: Final = next((r for r in tag_results if isinstance(r, BaseException)), None) + if tag_error is not None: + # Untag only the deployments this create actually appended to + # (add returned True). A deployment can already carry the tag + # name as a raw litellm_params string with no Tag row, and that + # pre-existing membership is not ours to remove. The fresh tag + # row still gets deleted, or a retry 400s on "already exists". + verbose_proxy_logger.exception("Rolling back tag %s after tagging failure: %s", tag.name, tag_error) + await _rollback_new_tag( + prisma_client=prisma_client, + tag_name=tag.name, + tagged_deployments=[d for d, r in zip(task_deployments, tag_results) if r is True], + # Only a budget this create minted is rolled back; a + # caller-supplied budget_id names a pre-existing row. + created_budget_id=budget_id if tag.budget_id is None else None, ) - await asyncio.gather(*tasks) + if isinstance(tag_error, HTTPException): + raise tag_error + raise HTTPException(status_code=500, detail=str(tag_error)) # Build response tag_config: Final = TagConfig( @@ -350,13 +380,19 @@ async def new_tag( "message": f"Tag {tag.name} created successfully", "tag": tag_config, } + except HTTPException: + # Deliberate raises above (the 400 for an existing tag, the rollback + # path's re-raise) reach the caller with their own status and detail. + raise except Exception as e: verbose_proxy_logger.exception("Error creating tag: %s", e) raise HTTPException(status_code=500, detail=str(e)) -async def _add_tag_to_deployment(deployment: "Deployment", tag: str): - """Helper function to add tag to deployment""" +async def _add_tag_to_deployment(deployment: "Deployment", tag: str) -> bool: + """Add a tag to one deployment. Returns True only when this call appended + it - False means the deployment already carried the tag (e.g. from config), + which /tag/new's rollback must leave alone.""" from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -383,8 +419,9 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): # Add tag to tags array (preserve encryption of other fields) if "tags" not in existing_params: existing_params["tags"] = [] - if tag not in existing_params["tags"]: - existing_params["tags"].append(tag) + if tag in existing_params["tags"]: + return False + existing_params["tags"].append(tag) # Update database with modified params (keeps encrypted fields encrypted) await ModelRepository(prisma_client).table.update( @@ -394,6 +431,72 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): except Exception as e: verbose_proxy_logger.exception("Error adding tag to deployment: %s", e) raise HTTPException(status_code=500, detail=str(e)) + else: + return True + + +async def _remove_tag_from_deployment(deployment: "Deployment", tag: str) -> None: + """Best-effort removal of a tag from a deployment's litellm_params.tags. + + Used only to roll back a failed /tag/new: the tag is brand new then, so + removing it cannot touch a tag anyone else created. Preserves encryption of + the other fields exactly like _add_tag_to_deployment. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return + + db_model: Final = await ModelRepository(prisma_client).table.find_unique( + where={"model_id": deployment.model_info.id} + ) + if db_model is None: + return + + raw_params: Final = db_model.litellm_params + existing_params: Final = cast( # cast-ok: prisma stores litellm_params as untyped JSON + "dict[str, object]", + json.loads(raw_params) if isinstance(raw_params, str) else raw_params, # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub is str + ) + if not isinstance(existing_params, dict): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub + return + + tags: Final = cast( # cast-ok: prisma stores litellm_params as untyped JSON + "list[object]", existing_params.get("tags") or [] + ) + if tag not in tags: + return + existing_params["tags"] = [t for t in tags if t != tag] + await ModelRepository(prisma_client).table.update( + where={"model_id": deployment.model_info.id}, + data={"litellm_params": json.dumps(existing_params)}, + ) + + +async def _rollback_new_tag( + prisma_client: "PrismaClient", + tag_name: str, + tagged_deployments: Sequence["Deployment"], + created_budget_id: str | None = None, +) -> None: + """Undo a failed /tag/new: untag only the deployments the create appended + the tag to (pre-existing raw-string membership stays), then delete the + orphaned tag row and any budget row the create minted. + Best-effort: the original tagging error is what the caller gets even if a + rollback step fails too.""" + try: + await asyncio.gather( + *(_remove_tag_from_deployment(deployment=deployment, tag=tag_name) for deployment in tagged_deployments), + return_exceptions=True, + ) + await _table(TagRepository(prisma_client)).delete(where={"tag_name": tag_name}) + if created_budget_id is not None: + await BudgetRepository(prisma_client).table.delete(where={"budget_id": created_budget_id}) + await _evict_tag_cache_keys((tag_cache_key(tag_name), tag_registry_cache_key())) + except Exception: # noqa: BLE001 - rollback is best-effort; the caller keeps the original tagging error + # CodeQL log-injection: tag_name is user-provided; strip line breaks before logging + safe_tag_name: Final = tag_name.replace("\r", " ").replace("\n", " ") + verbose_proxy_logger.exception("Rollback of failed tag create %s itself failed", safe_tag_name) @router.post( diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index e5d13a381ae..a73adc811ea 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -1536,12 +1536,13 @@ async def test_add_tag_to_deployment_no_duplicate_tags(): model_info=ModelInfo(id="model-789"), ) - # Try to add an existing tag - await _add_tag_to_deployment(deployment, "existing-tag") + # Try to add an existing tag - it is already there, so the helper + # reports False and skips the write (rollback relies on this contract) + result = await _add_tag_to_deployment(deployment, "existing-tag") - # Verify update was called - update_call = mock_db.litellm_proxymodeltable.update.call_args - updated_params = json.loads(update_call[1]["data"]["litellm_params"]) + assert result is False + mock_db.litellm_proxymodeltable.update.assert_not_called() + updated_params = db_model.litellm_params # Verify no duplicate tags assert updated_params["tags"].count("existing-tag") == 1 diff --git a/tests/unit/proxy/management_endpoints/test_tag_new_rollback.py b/tests/unit/proxy/management_endpoints/test_tag_new_rollback.py new file mode 100644 index 00000000000..bec62774cc7 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_tag_new_rollback.py @@ -0,0 +1,255 @@ +"""Regression tests for #47005: /tag/new wrote the tag row before tagging its +models, so a failure mid-create 500'd and left an orphan row that blocked every +retry with "already exists". The create now rolls back: untag wherever the new +tag landed, then delete the row.""" + +import contextlib +import json +from datetime import datetime, timezone +from types import SimpleNamespace +from typing import Any, Dict, List +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +import litellm.proxy.management_endpoints.tag_management_endpoints as endpoints +import litellm.proxy.proxy_server as proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.tag_management import TagNewRequest + +TAG_NAME = "rollback-tag" +NOW = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +def _deployment(model_id: str) -> SimpleNamespace: + return SimpleNamespace(model_info=SimpleNamespace(id=model_id)) + + +def _make_table() -> SimpleNamespace: + async def _create(data: Dict[str, Any]) -> SimpleNamespace: + return SimpleNamespace( + tag_name=data["tag_name"], + description=data["description"], + models=data["models"], + model_info=data["model_info"], + created_at=NOW, + updated_at=NOW, + created_by=data["created_by"], + ) + + return SimpleNamespace( + find_unique=AsyncMock(return_value=None), + create=AsyncMock(side_effect=_create), + delete=AsyncMock(), + ) + + +@contextlib.contextmanager +def _patched(table: SimpleNamespace, deployments: List[SimpleNamespace], add_side_effect): + add_tag: AsyncMock = AsyncMock(side_effect=add_side_effect) + remove_tag: AsyncMock = AsyncMock() + with ( + patch.object(endpoints, "_table", return_value=table), + patch.object(endpoints, "handle_budget_for_entity", AsyncMock(return_value=None)), + patch.object(endpoints, "_get_model_names", AsyncMock(return_value={})), + patch.object(endpoints, "_evict_tag_cache_keys", AsyncMock()), + patch.object(endpoints, "get_deployments_by_model", AsyncMock(return_value=deployments)), + patch.object(endpoints, "_add_tag_to_deployment", add_tag), + patch.object(endpoints, "_remove_tag_from_deployment", remove_tag), + patch.object(proxy_server, "prisma_client", MagicMock()), + patch.object(proxy_server, "llm_router", MagicMock()), + patch.object(proxy_server, "litellm_proxy_admin_name", "admin"), + ): + yield add_tag, remove_tag + + +@pytest.mark.asyncio +async def test_new_tag_rolls_back_row_and_deployments_when_tagging_fails(): + """One deployment write fails -> 500 to the caller, the tag is removed from + every deployment it reached, and the fresh row is deleted so a retry works.""" + table = _make_table() + deployments = [_deployment("model-1"), _deployment("model-2")] + with _patched(table, deployments, [True, HTTPException(status_code=500, detail="db write failed")]) as ( + add_tag, + remove_tag, + ): + with pytest.raises(HTTPException) as exc_info: + await endpoints.new_tag( + tag=TagNewRequest(name=TAG_NAME, models=["model-a"]), + user_api_key_dict=UserAPIKeyAuth(user_id="admin"), + ) + + assert exc_info.value.status_code == 500 + # the outer handler must not re-wrap the re-raise into a generic 500 + assert exc_info.value.detail == "db write failed" + # every deployment write was attempted (return_exceptions gather) + assert add_tag.await_count == 2 + # rollback: untag only the deployment this create actually appended to... + assert remove_tag.await_count == 1 + assert id(remove_tag.await_args_list[0].kwargs["deployment"]) == id(deployments[0]) + assert all(c.kwargs["tag"] == TAG_NAME for c in remove_tag.await_args_list) + # ...then delete the orphaned row + table.delete.assert_awaited_once_with(where={"tag_name": TAG_NAME}) + + +@pytest.mark.asyncio +async def test_new_tag_success_path_does_not_rollback(): + table = _make_table() + deployments = [_deployment("model-1")] + with _patched(table, deployments, [True]) as (add_tag, remove_tag): + resp = await endpoints.new_tag( + tag=TagNewRequest(name=TAG_NAME, models=["model-a"]), + user_api_key_dict=UserAPIKeyAuth(user_id="admin"), + ) + + assert resp["message"] == f"Tag {TAG_NAME} created successfully" + add_tag.assert_awaited_once() + remove_tag.assert_not_awaited() + table.delete.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_remove_tag_from_deployment_removes_only_that_tag(): + table = SimpleNamespace( + find_unique=AsyncMock( + return_value=SimpleNamespace( + litellm_params=json.dumps({"tags": [TAG_NAME, "other-tag"], "api_key": "encrypted-value"}) + ) + ), + update=AsyncMock(), + ) + with ( + patch.object(endpoints, "ModelRepository", return_value=SimpleNamespace(table=table)), + patch.object(proxy_server, "prisma_client", MagicMock()), + ): + await endpoints._remove_tag_from_deployment(deployment=_deployment("model-1"), tag=TAG_NAME) + + table.update.assert_awaited_once() + written = json.loads(table.update.await_args.kwargs["data"]["litellm_params"]) + assert written["tags"] == ["other-tag"] + assert written["api_key"] == "encrypted-value" + + +@pytest.mark.asyncio +async def test_remove_tag_from_deployment_noops_when_tag_absent(): # test-quality-ok: the no-op contract is the absence of a DB write, so assert_not_awaited on table.update is the caller-observable assertion + table = SimpleNamespace( + find_unique=AsyncMock(return_value=SimpleNamespace(litellm_params=json.dumps({"tags": ["other-tag"]}))), + update=AsyncMock(), + ) + with ( + patch.object(endpoints, "ModelRepository", return_value=SimpleNamespace(table=table)), + patch.object(proxy_server, "prisma_client", MagicMock()), + ): + await endpoints._remove_tag_from_deployment(deployment=_deployment("model-1"), tag=TAG_NAME) + + table.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_tag_non_http_tagging_failure_keeps_original_detail(): + """A plain exception from a deployment write becomes a 500 whose detail is + the original error text, not "500: ..." from a second wrap.""" + table = _make_table() + deployments = [_deployment("model-1")] + with _patched(table, deployments, [ValueError("deployment write kaboom")]): + with pytest.raises(HTTPException) as exc_info: + await endpoints.new_tag( + tag=TagNewRequest(name=TAG_NAME, models=["model-a"]), + user_api_key_dict=UserAPIKeyAuth(user_id="admin"), + ) + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "deployment write kaboom" + table.delete.assert_awaited_once_with(where={"tag_name": TAG_NAME}) + + +@pytest.mark.asyncio +async def test_new_tag_existing_tag_stays_a_400(): + """The pre-existing "already exists" raise must keep its 400 instead of + being turned into a 500 by the outer handler.""" + table = _make_table() + table.find_unique = AsyncMock(return_value=SimpleNamespace(tag_name=TAG_NAME)) + with _patched(table, [], []): + with pytest.raises(HTTPException) as exc_info: + await endpoints.new_tag( + tag=TagNewRequest(name=TAG_NAME, models=["model-a"]), + user_api_key_dict=UserAPIKeyAuth(user_id="admin"), + ) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == f"Tag {TAG_NAME} already exists" + + +@pytest.mark.asyncio +async def test_new_tag_rollback_deletes_budget_minted_by_the_create(): + """When the create minted a budget row (budget fields given, no budget_id), + the rollback deletes it so a retry does not leak stray budget rows.""" + table = _make_table() + deployments = [_deployment("model-1")] + budget_table = SimpleNamespace(delete=AsyncMock()) + with ( + _patched(table, deployments, [ValueError("boom")]), + patch.object(endpoints, "handle_budget_for_entity", AsyncMock(return_value="budget-minted-1")), + patch.object(endpoints, "BudgetRepository", return_value=SimpleNamespace(table=budget_table)), + ): + with pytest.raises(HTTPException): + await endpoints.new_tag( + tag=TagNewRequest(name=TAG_NAME, models=["model-a"], max_budget=10.0), + user_api_key_dict=UserAPIKeyAuth(user_id="admin"), + ) + + table.delete.assert_awaited_once_with(where={"tag_name": TAG_NAME}) + budget_table.delete.assert_awaited_once_with(where={"budget_id": "budget-minted-1"}) + + +@pytest.mark.asyncio +async def test_new_tag_rollback_keeps_caller_supplied_budget(): + """A budget_id supplied in the request names a pre-existing row; the + rollback must not delete it.""" + table = _make_table() + deployments = [_deployment("model-1")] + budget_table = SimpleNamespace(delete=AsyncMock()) + with ( + _patched(table, deployments, [ValueError("boom")]), + patch.object(endpoints, "handle_budget_for_entity", AsyncMock(return_value="budget-preexisting-1")), + patch.object(endpoints, "BudgetRepository", return_value=SimpleNamespace(table=budget_table)), + ): + with pytest.raises(HTTPException): + await endpoints.new_tag( + tag=TagNewRequest(name=TAG_NAME, models=["model-a"], budget_id="budget-preexisting-1"), + user_api_key_dict=UserAPIKeyAuth(user_id="admin"), + ) + + table.delete.assert_awaited_once_with(where={"tag_name": TAG_NAME}) + budget_table.delete.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_tag_rollback_preserves_preexisting_membership(): + """A deployment can already carry the tag name as a raw litellm_params + string (config-defined, no Tag row). _add_tag_to_deployment returns False + for it, and the rollback must leave that membership alone while untagging + the deployment the create actually appended to.""" + table = _make_table() + preexisting = _deployment("model-1") + appended = _deployment("model-2") + failing = _deployment("model-3") + deployments = [preexisting, appended, failing] + with _patched(table, deployments, [False, True, HTTPException(status_code=500, detail="db write failed")]) as ( + add_tag, + remove_tag, + ): + with pytest.raises(HTTPException) as exc_info: + await endpoints.new_tag( + tag=TagNewRequest(name=TAG_NAME, models=["model-a"]), + user_api_key_dict=UserAPIKeyAuth(user_id="admin"), + ) + + assert exc_info.value.status_code == 500 + assert add_tag.await_count == 3 + # only the deployment this create appended to gets untagged + remove_tag.assert_awaited_once() + assert id(remove_tag.await_args_list[0].kwargs["deployment"]) == id(appended) + # the orphan row is still deleted so a retry works + table.delete.assert_awaited_once_with(where={"tag_name": TAG_NAME})