From f6eb18436ed1692c1f89cfe603aa99f66eb68f4c Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Tue, 23 Jun 2026 10:44:39 -0700
Subject: [PATCH] fix(proxy): keep budget_id in typed customer responses

Address review feedback on the Customer Management response_model typing.
Greptile flagged that response_model=LiteLLM_EndUserTable on /customer/new
and /customer/update silently drops fields the raw Prisma model_dump()
echoed. Checking the schema, budget_id is the only such scalar column that
was missing from the Pydantic model (created_at/updated_at/tpm_limit do not
exist on litellm_endusertable), so add budget_id to LiteLLM_EndUserTable.

This restores budget_id on new/update and also fixes the pre-existing gap
where /customer/info and /customer/list (already typed) dropped it, which
the UI Customer type expects. A regression test pins budget_id surviving the
response_model filter on /customer/update.

Also document UnblockUsersResponse.blocked_users via a Field description: it
holds the users that remain blocked after the call. The key name predates
this PR and is kept to avoid a backwards-incompatible rename on a beta route.
---
 litellm/models/end_user.py                    |  1 +
 .../customer_endpoints.py                     |  6 +++--
 .../test_customer_endpoints.py                | 27 +++++++++++++++++++
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/litellm/models/end_user.py b/litellm/models/end_user.py
index 15fd03ec2ca..9bf895b9447 100644
--- a/litellm/models/end_user.py
+++ b/litellm/models/end_user.py
@@ -21,6 +21,7 @@ class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase):
     spend: float = 0.0
     allowed_model_region: Optional[Literal["eu", "us"]] = None
     default_model: Optional[str] = None
+    budget_id: Optional[str] = None
     litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
     object_permission_id: Optional[str] = None
     object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
diff --git a/litellm/types/proxy/management_endpoints/customer_endpoints.py b/litellm/types/proxy/management_endpoints/customer_endpoints.py
index 3aba8d2d211..52410220509 100644
--- a/litellm/types/proxy/management_endpoints/customer_endpoints.py
+++ b/litellm/types/proxy/management_endpoints/customer_endpoints.py
@@ -1,6 +1,6 @@
 from typing import List
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from litellm.models.end_user import LiteLLM_EndUserTable
 
@@ -10,7 +10,9 @@ class BlockUsersResponse(BaseModel):
 
 
 class UnblockUsersResponse(BaseModel):
-    blocked_users: List[str]
+    blocked_users: List[str] = Field(
+        description="User IDs that remain blocked after this unblock call"
+    )
 
 
 class DeleteCustomersResponse(BaseModel):
diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py
index 5243d868a7c..b3ca18e56c8 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py
@@ -360,6 +360,33 @@ def test_customer_new_documented_in_openapi_schema():
     assert json_schema["$ref"].endswith("/LiteLLM_EndUserTable")
 
 
+def test_update_customer_response_preserves_budget_id(
+    mock_prisma_client, mock_user_api_key_auth
+):
+    """
+    Regression for the response_model field-stripping concern: budget_id is a real
+    column on the end-user table that /customer/update echoes. response_model=
+    LiteLLM_EndUserTable must NOT drop it, so budget_id stays in LiteLLM_EndUserTable.
+    """
+    existing = LiteLLM_EndUserTable(user_id="cust-1", blocked=False)
+    updated = LiteLLM_EndUserTable(
+        user_id="cust-1", blocked=False, budget_id="budget-123"
+    )
+    mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(
+        return_value=existing
+    )
+    mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=updated)
+
+    response = client.post(
+        "/customer/update",
+        json={"user_id": "cust-1", "budget_id": "budget-123"},
+        headers={"Authorization": "Bearer test-key"},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["budget_id"] == "budget-123"
+
+
 def test_block_customer_success_serializes_through_response_model(
     mock_prisma_client, mock_user_api_key_auth
 ):