fix(proxy): return 404 for missing tags and preserve fields on partial tag update

litellm/proxy/management_endpoints/tag_management_endpoints.py

@@ -17,7 +17,7 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Final, Protocol, TypedDict, overload
 
 from fastapi import APIRouter, Depends, HTTPException, Query
-from pydantic import TypeAdapter
+from pydantic import TypeAdapter, ValidationError
 
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view
@@ -60,6 +60,7 @@ if TYPE_CHECKING:
 router: Final = APIRouter()
 
 _DECODED_JSON: Final = TypeAdapter(object)
+_MODEL_INFO_ADAPTER: Final = TypeAdapter(dict[str, str])
 
 
 class _TagRecord(Protocol):
@@ -221,6 +222,18 @@ async def _get_model_names(prisma_client: "PrismaClient", model_ids: Sequence[st
         return {}
 
 
+def _parse_model_info(raw_model_info: object) -> dict[str, str]:  # mutable-ok: TagConfig model_info requires dict
+    """Helper function to parse stored model_info into a dict[str, str]."""
+    if not raw_model_info:
+        return {}
+    try:
+        if isinstance(raw_model_info, str):
+            return _MODEL_INFO_ADAPTER.validate_json(raw_model_info)
+        return _MODEL_INFO_ADAPTER.validate_python(raw_model_info)
+    except (ValidationError, ValueError):
+        return {}
+
+
 async def get_deployments_by_model(model: str, llm_router: "Router") -> list["Deployment"]:
     """
     Get all deployments by model
@@ -348,6 +361,8 @@ async def new_tag(
             "message": f"Tag {tag.name} created successfully",
             "tag": tag_config,
         }
+    except HTTPException:
+        raise
     except Exception as e:
         verbose_proxy_logger.exception("Error creating tag: %s", e)
         raise HTTPException(status_code=500, detail=str(e))
@@ -444,15 +459,20 @@ async def update_tag(
             cleared_budget_fields=frozenset(field for field in tag.model_fields_set if getattr(tag, field) is None),
         )
 
-        # Get model names for model_info
-        model_info: Final = await _get_model_names(prisma_client, tag.models or [])
-
         # Prepare update data
-        update_data: Final = {
-            "description": tag.description,
-            "models": tag.models or [],
-            "model_info": json.dumps(model_info),
-        }
+        update_data: Final[dict[str, object]] = {}  # mutable-ok: prisma update input payload built conditionally
+
+        if "description" in tag.model_fields_set:
+            update_data["description"] = tag.description
+
+        model_info: Final = (
+            await _get_model_names(prisma_client, tag.models or [])
+            if "models" in tag.model_fields_set
+            else _parse_model_info(getattr(existing_tag, "model_info", None))
+        )
+        if "models" in tag.model_fields_set:
+            update_data["models"] = tag.models or []
+            update_data["model_info"] = json.dumps(model_info)
 
         # Add budget_id if it changed
         if budget_id != existing_tag.budget_id:
@@ -481,6 +501,8 @@ async def update_tag(
             "message": f"Tag {tag.name} updated successfully",
             "tag": tag_config,
         }
+    except HTTPException:
+        raise
     except Exception as e:
         verbose_proxy_logger.exception("Error updating tag: %s", e)
         raise HTTPException(status_code=500, detail=str(e))
@@ -547,6 +569,8 @@ async def info_tag(
             requested_tags[tag_record.tag_name] = tag_dict
 
         return requested_tags
+    except HTTPException:
+        raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -718,6 +742,8 @@ async def delete_tag(
         await _evict_tag_cache_keys((tag_cache_key(data.name), tag_registry_cache_key()))
 
         return {"message": f"Tag {data.name} deleted successfully"}
+    except HTTPException:
+        raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 

tests/integration/spend/test_tag_budget_enforcement.py

@@ -492,8 +492,6 @@ def test_tag_delete_evicts_blocked_limit(
 
 
 def test_tag_info_returns_404_for_deleted_tag(gateway: Gateway) -> None:
-    pytest.skip("BUG: POST /tag/info for a deleted tag returns 500 with detail '404: Tags not found' instead of 404")
-
     tag: Final = f"tag-info-deleted-{uuid.uuid4().hex}"
     with gateway.scenario() as scenario:
 
@@ -511,8 +509,6 @@ def test_tag_info_returns_404_for_deleted_tag(gateway: Gateway) -> None:
 
 
 def test_tag_partial_update_keeps_description_and_models(gateway: Gateway) -> None:
-    pytest.skip("BUG: /tag/update without description/models wipes the stored description and models")
-
     tag: Final = f"tag-partial-update-{uuid.uuid4().hex}"
     with gateway.scenario() as scenario:
         model: Final = scenario.model()

tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py

@@ -1586,3 +1586,211 @@ def test_tag_info_and_tag_list_return_the_stored_model_info_decoded(
     assert info_response.json()["routed-tag"]["model_info"] == returned_model_info
     assert list_response.status_code == 200
     assert [tag["model_info"] for tag in list_response.json()] == [returned_model_info]
+
+
+def test_tag_info_returns_404_when_tag_not_found(monkeypatch):
+    from litellm.proxy import proxy_server
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+    mock_db = Mock()
+    mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[])
+
+    monkeypatch.setattr(
+        proxy_server,
+        "prisma_client",
+        SimpleNamespace(db=mock_db),
+    )
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    try:
+        response = client.post("/tag/info", json={"names": ["nonexistent-tag"]})
+    finally:
+        app.dependency_overrides.clear()
+
+    assert response.status_code == 404
+    assert response.json()["detail"] == "Tags not found: ['nonexistent-tag']"
+
+
+def test_tag_update_returns_404_when_tag_not_found(monkeypatch):
+    from litellm.proxy import proxy_server
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+    mock_db = Mock()
+    mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=None)
+
+    monkeypatch.setattr(
+        proxy_server,
+        "prisma_client",
+        SimpleNamespace(db=mock_db),
+    )
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    try:
+        response = client.post("/tag/update", json={"name": "nonexistent-tag", "description": "new"})
+    finally:
+        app.dependency_overrides.clear()
+
+    assert response.status_code == 404
+    assert "not found" in response.json()["detail"]
+
+
+def test_tag_delete_returns_404_when_tag_not_found(monkeypatch):
+    from litellm.proxy import proxy_server
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+    mock_db = Mock()
+    mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=None)
+
+    monkeypatch.setattr(
+        proxy_server,
+        "prisma_client",
+        SimpleNamespace(db=mock_db),
+    )
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    try:
+        response = client.post("/tag/delete", json={"name": "nonexistent-tag"})
+    finally:
+        app.dependency_overrides.clear()
+
+    assert response.status_code == 404
+    assert "not found" in response.json()["detail"]
+
+
+def test_tag_partial_update_does_not_wipe_description_or_models(monkeypatch):
+    from litellm.proxy import proxy_server
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+    existing_tag = SimpleNamespace(
+        tag_name="test-tag",
+        description="existing description",
+        models=["model-1"],
+        model_info='{"model-1": "gpt-4"}',
+        budget_id=None,
+        created_at=datetime(2025, 1, 1),
+        updated_at=datetime(2025, 1, 2),
+        created_by="user-123",
+    )
+    updated_tag = SimpleNamespace(
+        tag_name="test-tag",
+        description="existing description",
+        models=["model-1"],
+        model_info='{"model-1": "gpt-4"}',
+        budget_id="budget-1",
+        created_at=datetime(2025, 1, 1),
+        updated_at=datetime(2025, 1, 3),
+        created_by="user-123",
+    )
+
+    mock_db = Mock()
+    mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag)
+    captured_update_kwargs: dict[str, object] = {}
+
+    async def fake_update(**kwargs):
+        captured_update_kwargs.update(kwargs)
+        return updated_tag
+
+    mock_db.litellm_tagtable.update = fake_update
+
+    budget_record = SimpleNamespace(budget_id="budget-1", max_budget=10.0)
+    mock_db.litellm_budgettable.create = AsyncMock(return_value=budget_record)
+
+    monkeypatch.setattr(
+        proxy_server,
+        "prisma_client",
+        SimpleNamespace(db=mock_db, jsonify_object=lambda data: dict(data)),
+    )
+    monkeypatch.setattr(proxy_server, "litellm_proxy_admin_name", "admin")
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    try:
+        response = client.post("/tag/update", json={"name": "test-tag", "max_budget": 10.0})
+    finally:
+        app.dependency_overrides.clear()
+
+    assert response.status_code == 200
+    res_data = response.json()
+    assert res_data["tag"]["description"] == "existing description"
+    assert res_data["tag"]["models"] == ["model-1"]
+    assert res_data["tag"]["model_info"] == {"model-1": "gpt-4"}
+    data_sent = captured_update_kwargs.get("data", {})
+    assert "description" not in data_sent
+    assert "models" not in data_sent
+    assert "model_info" not in data_sent
+
+
+def test_new_tag_duplicate_returns_400(monkeypatch):
+    from litellm.proxy import proxy_server
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+    mock_db = Mock()
+    existing_tag = SimpleNamespace(tag_name="existing-tag")
+    mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag)
+
+    monkeypatch.setattr(
+        proxy_server,
+        "prisma_client",
+        SimpleNamespace(db=mock_db),
+    )
+    monkeypatch.setattr(proxy_server, "llm_router", Mock())
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    try:
+        response = client.post("/tag/new", json={"name": "existing-tag"})
+    finally:
+        app.dependency_overrides.clear()
+
+    assert response.status_code == 400
+    assert "already exists" in response.json()["detail"]
+
+
+def test_tag_update_with_invalid_model_info_returns_empty_dict(monkeypatch):
+    from litellm.proxy import proxy_server
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+    existing_tag = SimpleNamespace(
+        tag_name="test-tag",
+        description="existing",
+        models=["model-1"],
+        model_info="invalid-json-data",
+        budget_id=None,
+        created_at=datetime(2025, 1, 1),
+        updated_at=datetime(2025, 1, 2),
+        created_by="user-123",
+    )
+    updated_tag = SimpleNamespace(
+        tag_name="test-tag",
+        description="existing",
+        models=["model-1"],
+        model_info="invalid-json-data",
+        budget_id=None,
+        created_at=datetime(2025, 1, 1),
+        updated_at=datetime(2025, 1, 3),
+        created_by="user-123",
+    )
+
+    mock_db = Mock()
+    mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag)
+    mock_db.litellm_tagtable.update = AsyncMock(return_value=updated_tag)
+
+    monkeypatch.setattr(
+        proxy_server,
+        "prisma_client",
+        SimpleNamespace(db=mock_db, jsonify_object=lambda data: dict(data)),
+    )
+    monkeypatch.setattr(proxy_server, "litellm_proxy_admin_name", "admin")
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    try:
+        response = client.post("/tag/update", json={"name": "test-tag", "description": "new-desc"})
+    finally:
+        app.dependency_overrides.clear()
+
+    assert response.status_code == 200
+    assert response.json()["tag"]["model_info"] == {}