Merge 0687d145f5 into c905cd065d

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(

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

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})