Merge remote-tracking branch 'origin/litellm_team_member_temp_budget_increase' into litellm_team_member_temp_budget_increase

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	tests/test_litellm/proxy/auth/test_auth_checks.py

litellm/proxy/management_endpoints/organization_endpoints.py

@@ -291,6 +291,9 @@ async def _verify_org_access(
 _STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
 _BUDGET_SETTABLE_FIELDS: Final = frozenset(LiteLLM_BudgetTable.model_fields.keys()) - {"budget_id"}
 _ORG_COLUMN_FIELDS: Final = frozenset({"organization_alias", "models"})
+_ORG_METADATA_FIELDS: Final = tuple(
+    field for field in LiteLLM_ManagementEndpoint_MetadataFields if field not in _BUDGET_SETTABLE_FIELDS
+)
 
 
 def build_budget_write_data(budget_updates: Mapping[str, object], updated_by: str) -> Mapping[str, object]:
@@ -514,7 +517,7 @@ async def new_organization(
     organization_payload["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
     organization_row: Final = LiteLLM_OrganizationTable.model_validate(organization_payload)
 
-    for field in LiteLLM_ManagementEndpoint_MetadataFields:
+    for field in _ORG_METADATA_FIELDS:
         if getattr(data, field, None) is not None:
             _set_object_metadata_field(
                 object_data=organization_row,

tests/test_litellm/models/test_models.py

@@ -80,12 +80,10 @@ class TestBudget:
         assert budget.effective_max_budget(now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 150.0
 
     def test_effective_max_budget_ignores_expired_increase(self):
-        budget = LiteLLM_BudgetTable(
-            max_budget=100.0,
-            temp_budget_increase=50.0,
-            temp_budget_expiry=datetime(2020, 1, 1, tzinfo=timezone.utc),
-        )
+        expiry = datetime(2020, 1, 1, tzinfo=timezone.utc)
+        budget = LiteLLM_BudgetTable(max_budget=100.0, temp_budget_increase=50.0, temp_budget_expiry=expiry)
         assert budget.effective_max_budget(now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 100.0
+        assert budget.effective_max_budget(now=expiry) == 100.0
 
     def test_effective_max_budget_without_increase(self):
         now = datetime(2026, 1, 1, tzinfo=timezone.utc)

tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py

@@ -1346,6 +1346,49 @@ async def test_new_organization_rejects_shared_alias_tool_permission_key():
     prisma_client.db.litellm_objectpermissiontable.create.assert_not_called()
 
 
+@pytest.mark.asyncio
+async def test_new_organization_temp_budget_fields_go_to_budget_row_not_metadata(monkeypatch):
+    """temp_budget_increase/expiry are budget columns and also key-metadata field names, so
+    /organization/new must write them to the budget row and keep the datetime out of the org
+    metadata JSON (a datetime there broke JSON serialization and 500'd the request)."""
+    from datetime import datetime, timezone
+
+    from litellm.proxy._types import LitellmUserRoles, NewOrganizationRequest, UserAPIKeyAuth
+    from litellm.proxy.management_endpoints.organization_endpoints import new_organization
+    from litellm.proxy.utils import PrismaClient
+
+    expiry = datetime(2099, 1, 1, tzinfo=timezone.utc)
+    prisma_client = MagicMock()
+    prisma_client.jsonify_object = MagicMock(side_effect=lambda data: PrismaClient.jsonify_object(prisma_client, data))
+    prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
+    prisma_client.db.litellm_budgettable.create = AsyncMock(return_value=MagicMock(budget_id="budget-1"))
+    prisma_client.db.litellm_organizationtable.create = AsyncMock(return_value={"organization_id": "org-1"})
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False)
+
+    response = await new_organization(
+        data=NewOrganizationRequest(
+            organization_alias="org",
+            max_budget=10,
+            temp_budget_increase=5,
+            temp_budget_expiry=expiry,
+        ),
+        user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
+    )
+
+    assert response == {"organization_id": "org-1"}
+    budget_write = prisma_client.db.litellm_budgettable.create.await_args.kwargs["data"]
+    assert (budget_write["max_budget"], budget_write["temp_budget_increase"], budget_write["temp_budget_expiry"]) == (
+        10,
+        5,
+        expiry,
+    )
+    org_write = prisma_client.db.litellm_organizationtable.create.await_args.kwargs["data"]
+    assert org_write["budget_id"] == "budget-1"
+    assert json.loads(org_write.get("metadata", "{}")) == {}
+
+
 def test_v2_update_organization_is_in_openapi_schema():
     """PATCH /v2/organization/{organization_id} is documented in the generated OpenAPI spec."""
     from fastapi import FastAPI