feat(keys): allow editing soft budget on existing keys (#39002)

* feat(keys): allow editing soft budget on existing keys

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

* refactor(ui): extract KeyBudgetNumberField to keep key_edit_view under max-lines

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

* style(ui): format keyEditFormValues with prettier

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

* test(keys): cover soft budget validation and update adapter

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

* fix(ui): reject non-finite soft budget values instead of clearing

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

* test(keys): assert soft budget validation returns None for valid values

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

* fix(keys): write soft budget and key row in one transaction

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

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yassin <yassin@berri.ai>

litellm/proxy/_types.py

@@ -1276,6 +1276,7 @@ class UpdateKeyRequest(KeyRequestBase):
     # else they will get overwritten
     duration: str | None = None
     spend: float | None = None
+    soft_budget: float | None = None
     metadata: dict | None = None
     temp_budget_increase: float | None = None
     temp_budget_expiry: datetime | None = None

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -196,6 +196,38 @@ class _ModelRowWhere(TypedDict):
     model_id: ReadOnly[str]
 
 
+class _KeyUpdateResult(TypedDict):
+    token: ReadOnly[str]
+    data: ReadOnly[Mapping[str, object]]
+
+
+class _KeyRowWhere(TypedDict):
+    token: ReadOnly[str]
+
+
+class _BudgetRowWhere(TypedDict):
+    budget_id: ReadOnly[str]
+
+
+class _BudgetRowSoftBudgetUpdate(TypedDict):
+    soft_budget: ReadOnly[float | None]
+    updated_by: ReadOnly[str]
+
+
+class _BudgetRowSoftBudgetCreate(TypedDict):
+    soft_budget: ReadOnly[float]
+    created_by: ReadOnly[str]
+    updated_by: ReadOnly[str]
+
+
+class _KeyUpdateTx(Protocol):
+    @property
+    def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ...
+
+    @property
+    def litellm_budgettable(self) -> "TableActions[prisma_models.LiteLLM_BudgetTable]": ...
+
+
 class _ConfigTableActions(Protocol):
     """Config table surface this module needs; the shared repository seam exposes no ``update``."""
 
@@ -1812,11 +1844,7 @@ async def generate_key_fn(
                 status_code=400,
                 detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"},
             )
-        if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0):
-            raise HTTPException(
-                status_code=400,
-                detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
-            )
+        _validate_soft_budget_value(data.soft_budget)
 
         custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = (
             _custom_key_generate_hook(proxy_server)
@@ -2121,6 +2149,88 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_
     return non_default_values
 
 
+def _validate_soft_budget_value(soft_budget: float | None) -> None:
+    if soft_budget is not None and (not math.isfinite(soft_budget) or soft_budget < 0):
+        raise HTTPException(
+            status_code=400,
+            detail={"error": f"soft_budget must be a non-negative finite number. Received: {soft_budget}"},
+        )
+
+
+async def _update_key_soft_budget(
+    db: _KeyUpdateTx,
+    existing_key_row: LiteLLM_VerificationToken,
+    soft_budget: float | None,
+    changed_by: str,
+) -> str | None:
+    existing_budget_id: Final = existing_key_row.budget_id
+    if existing_budget_id is not None:
+        budget_update: Final[_BudgetRowSoftBudgetUpdate] = {"soft_budget": soft_budget, "updated_by": changed_by}
+        budget_where: Final[_BudgetRowWhere] = {"budget_id": existing_budget_id}
+        await db.litellm_budgettable.update(where=budget_where, data=budget_update)
+        return existing_budget_id
+    if soft_budget is None:
+        return None
+    budget_create: Final[_BudgetRowSoftBudgetCreate] = {
+        "soft_budget": soft_budget,
+        "created_by": changed_by,
+        "updated_by": changed_by,
+    }
+    created_budget: Final = await db.litellm_budgettable.create(data=budget_create)
+    return created_budget.budget_id
+
+
+async def _apply_soft_budget_update(
+    data: UpdateKeyRequest,
+    non_default_values: Mapping[str, object],
+    db: _KeyUpdateTx,
+    existing_key_row: LiteLLM_VerificationToken,
+    changed_by: str,
+) -> Mapping[str, object]:
+    remaining: Final = MappingProxyType({k: v for k, v in non_default_values.items() if k != "soft_budget"})
+    updated_budget_id: Final = await _update_key_soft_budget(
+        db=db,
+        existing_key_row=existing_key_row,
+        soft_budget=data.soft_budget,
+        changed_by=changed_by,
+    )
+    if updated_budget_id is not None and existing_key_row.budget_id is None:
+        return MappingProxyType({**remaining, "budget_id": updated_budget_id})
+    return remaining
+
+
+async def _update_key_row_with_soft_budget(
+    prisma_client: PrismaClient,
+    key: str,
+    data: UpdateKeyRequest,
+    non_default_values: Mapping[str, object],
+    existing_key_row: LiteLLM_VerificationToken,
+    changed_by: str,
+) -> _KeyUpdateResult:
+    hashed_token: Final = _hash_token_if_needed(key)
+    key_where: Final[_KeyRowWhere] = {"token": hashed_token}
+    tx: _KeyUpdateTx
+    async with prisma_client.tx() as tx:
+        update_values: Final = await _apply_soft_budget_update(
+            data=data,
+            non_default_values=non_default_values,
+            db=tx,
+            existing_key_row=existing_key_row,
+            changed_by=changed_by,
+        )
+        updated_row: Final = await tx.litellm_verificationtoken.update(
+            where=key_where,
+            data=with_settings_updated_at(
+                prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token}))
+            ),
+        )
+    updated_data: Final[Mapping[str, object]] = (
+        updated_row.model_dump() if updated_row is not None else MappingProxyType({})
+    )
+    result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data}
+    return result
+
+
 async def prepare_key_update_data(
     data: UpdateKeyRequest | RegenerateKeyRequest,
     existing_key_row: LiteLLM_VerificationToken,
@@ -2659,6 +2769,7 @@ async def _validate_update_key_data(
         (data.max_budget is not None and data.max_budget != existing_key_row.max_budget)
         or data.spend is not None
         or "budget_limits" in data.model_fields_set
+        or "soft_budget" in data.model_fields_set
     )
 
     _existing_metadata: Final = getattr(existing_key_row, "metadata", None)
@@ -2862,7 +2973,7 @@ async def update_key_fn(
     - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
     - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
     - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
-    - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
+    - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget.
     - max_parallel_requests: Optional[int] - Rate limit for parallel requests
     - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"}
     - tpm_limit: Optional[int] - Tokens per minute limit
@@ -2918,6 +3029,7 @@ async def update_key_fn(
     """
     from litellm.proxy import proxy_server
     from litellm.proxy.proxy_server import (
+        litellm_proxy_admin_name,
         llm_router,
         premium_user,
         prisma_client,
@@ -2933,6 +3045,8 @@ async def update_key_fn(
                 detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"},
             )
 
+        _validate_soft_budget_value(data.soft_budget)
+
         # get the row from db
         existing_key_row: Final = await _get_and_validate_existing_key(
             token=data.key,
@@ -2989,10 +3103,22 @@ async def update_key_fn(
             existing_key_alias=existing_key_row.key_alias,
         )
 
-        _data: Final = {**non_default_values, "token": key}
         if prisma_client is None:
             raise Exception("Not connected to DB!")
-        response: Final = await prisma_client.update_data(token=key, data=_data)
+
+        changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name
+        response: Final = (
+            await _update_key_row_with_soft_budget(
+                prisma_client=prisma_client,
+                key=key,
+                data=data,
+                non_default_values=non_default_values,
+                existing_key_row=existing_key_row,
+                changed_by=changed_by,
+            )
+            if "soft_budget" in data.model_fields_set
+            else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key}))
+        )
 
         # Delete - key from cache, since it's been updated!
         # key updated - a new model could have been added to this key. it should not block requests after this is done
@@ -6432,7 +6558,7 @@ async def _list_key_helper(
                     {"token": "desc"},  # fallback sort
                 ]
             ),
-            include={"object_permission": True},
+            include={"object_permission": True, "litellm_budget_table": True},
         )
 
     verbose_proxy_logger.debug("Fetched %s keys", len(keys))

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -17692,6 +17692,253 @@ async def test_check_project_key_limits_still_rejects_real_model_outside_project
     assert "Model 'gpt-5.4-mini' not in project's allowed models" in exc_info.value.detail["error"]
 
 
+@pytest.mark.asyncio
+async def test_update_key_soft_budget_updates_existing_budget_row():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123")
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.update = AsyncMock()
+    mock_db.litellm_budgettable.create = AsyncMock()
+
+    result = await _update_key_soft_budget(
+        db=mock_db,
+        existing_key_row=existing_key,
+        soft_budget=25.0,
+        changed_by="user-1",
+    )
+
+    assert result == "budget-123"
+    mock_db.litellm_budgettable.update.assert_awaited_once_with(
+        where={"budget_id": "budget-123"},
+        data={"soft_budget": 25.0, "updated_by": "user-1"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_key_soft_budget_clears_existing_budget_row():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123")
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.update = AsyncMock()
+    mock_db.litellm_budgettable.create = AsyncMock()
+
+    result = await _update_key_soft_budget(
+        db=mock_db,
+        existing_key_row=existing_key,
+        soft_budget=None,
+        changed_by="user-1",
+    )
+
+    assert result == "budget-123"
+    mock_db.litellm_budgettable.update.assert_awaited_once_with(
+        where={"budget_id": "budget-123"},
+        data={"soft_budget": None, "updated_by": "user-1"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_key_soft_budget_creates_budget_row_when_key_has_none():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    created_row = MagicMock()
+    created_row.budget_id = "budget-new"
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.create = AsyncMock(return_value=created_row)
+    mock_db.litellm_budgettable.update = AsyncMock()
+
+    result = await _update_key_soft_budget(
+        db=mock_db,
+        existing_key_row=existing_key,
+        soft_budget=10.5,
+        changed_by="user-1",
+    )
+
+    assert result == "budget-new"
+    mock_db.litellm_budgettable.create.assert_awaited_once_with(
+        data={"soft_budget": 10.5, "created_by": "user-1", "updated_by": "user-1"}
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_key_soft_budget_noop_when_clearing_without_budget_row():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.create = AsyncMock()
+    mock_db.litellm_budgettable.update = AsyncMock()
+
+    result = await _update_key_soft_budget(
+        db=mock_db,
+        existing_key_row=existing_key,
+        soft_budget=None,
+        changed_by="user-1",
+    )
+
+    assert result is None
+    mock_db.litellm_budgettable.create.assert_not_awaited()
+    mock_db.litellm_budgettable.update.assert_not_awaited()
+
+
+def test_update_key_request_accepts_soft_budget():
+    request = UpdateKeyRequest(key="sk-test", soft_budget=42.0)
+    assert request.soft_budget == 42.0
+    assert "soft_budget" in request.model_fields_set
+
+
+@pytest.mark.parametrize("valid_value", [None, 0.0, 25.0])
+def test_validate_soft_budget_value_accepts_valid_values(valid_value):
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _validate_soft_budget_value,
+    )
+
+    assert _validate_soft_budget_value(valid_value) is None
+
+
+@pytest.mark.parametrize("invalid_value", [-5.0, float("nan"), float("inf")])
+def test_validate_soft_budget_value_rejects_invalid_values(invalid_value):
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _validate_soft_budget_value,
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        _validate_soft_budget_value(invalid_value)
+
+    assert exc_info.value.status_code == 400
+    assert "soft_budget must be a non-negative finite number" in str(exc_info.value.detail)
+
+
+@pytest.mark.asyncio
+async def test_apply_soft_budget_update_adds_budget_id_for_new_budget_row():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _apply_soft_budget_update,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    created_row = MagicMock()
+    created_row.budget_id = "budget-created-456"
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.create = AsyncMock(return_value=created_row)
+    mock_db.litellm_budgettable.update = AsyncMock()
+
+    result = await _apply_soft_budget_update(
+        data=UpdateKeyRequest(key="sk-test", soft_budget=25.0),
+        non_default_values={"soft_budget": 25.0},
+        db=mock_db,
+        existing_key_row=existing_key,
+        changed_by="user-1",
+    )
+
+    assert dict(result) == {"budget_id": "budget-created-456"}
+
+
+@pytest.mark.asyncio
+async def test_apply_soft_budget_update_keeps_existing_budget_id_out_of_token_update():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _apply_soft_budget_update,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123")
+    mock_db = MagicMock()
+    mock_db.litellm_budgettable.update = AsyncMock()
+    mock_db.litellm_budgettable.create = AsyncMock()
+
+    result = await _apply_soft_budget_update(
+        data=UpdateKeyRequest(key="sk-test", soft_budget=40.0),
+        non_default_values={"soft_budget": 40.0, "max_budget": 100.0},
+        db=mock_db,
+        existing_key_row=existing_key,
+        changed_by="user-1",
+    )
+
+    assert dict(result) == {"max_budget": 100.0}
+    mock_db.litellm_budgettable.update.assert_awaited_once_with(
+        where={"budget_id": "budget-123"},
+        data={"soft_budget": 40.0, "updated_by": "user-1"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_key_row_with_soft_budget_updates_budget_and_key_in_transaction():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_row_with_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    created_row = MagicMock(budget_id="budget-new")
+    updated_row = MagicMock()
+    updated_row.model_dump.return_value = {"token": "hashed", "budget_id": "budget-new"}
+    tx = MagicMock()
+    tx.litellm_budgettable.create = AsyncMock(return_value=created_row)
+    tx.litellm_verificationtoken.update = AsyncMock(return_value=updated_row)
+    tx_context = MagicMock()
+    tx_context.__aenter__ = AsyncMock(return_value=tx)
+    tx_context.__aexit__ = AsyncMock(return_value=None)
+    prisma_client = MagicMock()
+    prisma_client.tx.return_value = tx_context
+    prisma_client.jsonify_object = lambda data: dict(data)
+
+    result = await _update_key_row_with_soft_budget(
+        prisma_client=prisma_client,
+        key="sk-test",
+        data=UpdateKeyRequest(key="sk-test", soft_budget=25.0),
+        non_default_values={"soft_budget": 25.0},
+        existing_key_row=existing_key,
+        changed_by="user-1",
+    )
+
+    assert set(result) == {"token", "data"}
+    assert result["data"] == {"token": "hashed", "budget_id": "budget-new"}
+    tx.litellm_verificationtoken.update.assert_awaited_once()
+    update_call = tx.litellm_verificationtoken.update.await_args
+    assert update_call.kwargs["where"] == {"token": result["token"]}
+    assert update_call.kwargs["data"]["budget_id"] == "budget-new"
+    assert "soft_budget" not in update_call.kwargs["data"]
+
+
+@pytest.mark.asyncio
+async def test_update_key_row_with_soft_budget_propagates_transaction_error():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _update_key_row_with_soft_budget,
+    )
+
+    existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None)
+    created_row = MagicMock(budget_id="budget-new")
+    tx = MagicMock()
+    tx.litellm_budgettable.create = AsyncMock(return_value=created_row)
+    tx.litellm_verificationtoken.update = AsyncMock(side_effect=RuntimeError("update failed"))
+    tx_context = MagicMock()
+    tx_context.__aenter__ = AsyncMock(return_value=tx)
+    tx_context.__aexit__ = AsyncMock(return_value=None)
+    prisma_client = MagicMock()
+    prisma_client.tx.return_value = tx_context
+    prisma_client.jsonify_object = lambda data: dict(data)
+
+    with pytest.raises(RuntimeError, match="update failed"):
+        await _update_key_row_with_soft_budget(
+            prisma_client=prisma_client,
+            key="sk-test",
+            data=UpdateKeyRequest(key="sk-test", soft_budget=25.0),
+            non_default_values={"soft_budget": 25.0},
+            existing_key_row=existing_key,
+            changed_by="user-1",
+        )
+
+    tx_context.__aexit__.assert_awaited_once()
+    assert tx_context.__aexit__.await_args.args[0] is RuntimeError
+
+
 def test_generate_key_request_blank_team_id_is_personal():
     """The UI Team-field clear submits team_id=""; it must count as no team (LIT-3925)."""
     from litellm.proxy._types import RegenerateKeyRequest

ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx

@@ -1,7 +1,11 @@
 import React from "react";
+import { Control } from "react-hook-form";
 import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import { CircleHelp } from "lucide-react";
+import { FormField } from "@/components/shared/form/FormField";
+import NumericalInput from "../shared/numerical_input";
+import { KeyEditFormValues } from "./keyEditFormValues";
 
 export const labelWithHint = (label: React.ReactNode, hint: string): React.ReactNode => (
   <>
@@ -48,3 +52,27 @@ export const KeyTypeSelect = ({
     </SelectContent>
   </Select>
 );
+
+export const KeyBudgetNumberField = ({
+  control,
+  name,
+  label,
+  placeholder,
+}: {
+  control: Control<KeyEditFormValues>;
+  name: "max_budget" | "soft_budget";
+  label: string;
+  placeholder: string;
+}) => (
+  <FormField control={control} name={name} label={label}>
+    {({ ref: _ref, ...field }) => (
+      <NumericalInput
+        {...field}
+        value={field.value ?? ""}
+        step={0.01}
+        style={{ width: "100%" }}
+        placeholder={placeholder}
+      />
+    )}
+  </FormField>
+);

ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx

@@ -1,5 +1,10 @@
 import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { toast } from "@/lib/toast";
+
+vi.mock("@/lib/toast", () => ({
+  toast: { error: vi.fn(), success: vi.fn() },
+}));
 
 // ---- Hoisted shared mocks (safe to use inside vi.mock factories) ----
 const { keyUpdateCallMock, keyDeleteCallMock, mockUseAuthorized } = vi.hoisted(() => {
@@ -483,3 +488,104 @@ describe("KeyInfoView handleKeyUpdate empty strings", () => {
     });
   });
 });
+
+describe("KeyInfoView handleKeyUpdate soft_budget", () => {
+  const premiumAdminAuth = {
+    accessToken: "access_abc",
+    userId: "user_1",
+    userRole: "Admin",
+    premiumUser: true,
+    token: "token_123",
+    userEmail: "test@example.com",
+    disabledPersonalKeyCreation: false,
+    showSSOBanner: false,
+  };
+
+  const renderWithSoftBudget = (softBudget: number | null) => {
+    mockUseAuthorized.mockReturnValue(premiumAdminAuth);
+
+    return render(
+      <KeyInfoView
+        keyId="tok_123"
+        onClose={() => {}}
+        keyData={
+          { ...baseKeyData, litellm_budget_table: softBudget === null ? null : { soft_budget: softBudget } } as any
+        }
+        onKeyDataUpdate={() => {}}
+        teams={[]}
+      />,
+    );
+  };
+
+  it("should send a changed soft_budget as a number", async () => {
+    renderWithSoftBudget(null);
+
+    fireEvent.click(screen.getByText("Settings"));
+    fireEvent.click(screen.getByText("Edit Settings"));
+    (globalThis as any).__TEST_FORM_VALUES = {
+      token: "tok_123",
+      soft_budget: "25",
+    };
+
+    fireEvent.click(screen.getByText("Mock Submit"));
+
+    await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled());
+
+    const [, sentPayload] = keyUpdateCallMock.mock.calls[0];
+    expect(sentPayload.soft_budget).toBe(25);
+  });
+
+  it("should omit an unchanged soft_budget so unrelated edits skip the budget gate", async () => {
+    renderWithSoftBudget(25);
+
+    fireEvent.click(screen.getByText("Settings"));
+    fireEvent.click(screen.getByText("Edit Settings"));
+    (globalThis as any).__TEST_FORM_VALUES = {
+      token: "tok_123",
+      soft_budget: 25,
+      key_alias: "renamed",
+    };
+
+    fireEvent.click(screen.getByText("Mock Submit"));
+
+    await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled());
+
+    const [, sentPayload] = keyUpdateCallMock.mock.calls[0];
+    expect("soft_budget" in sentPayload).toBe(false);
+  });
+
+  it("should forward a cleared soft_budget as an explicit null the JSON body keeps", async () => {
+    renderWithSoftBudget(25);
+
+    fireEvent.click(screen.getByText("Settings"));
+    fireEvent.click(screen.getByText("Edit Settings"));
+    (globalThis as any).__TEST_FORM_VALUES = {
+      token: "tok_123",
+      soft_budget: "",
+    };
+
+    fireEvent.click(screen.getByText("Mock Submit"));
+
+    await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled());
+
+    const [, sentPayload] = keyUpdateCallMock.mock.calls[0];
+    expect(sentPayload.soft_budget).toBeNull();
+    expect(JSON.stringify({ ...sentPayload })).toContain('"soft_budget":null');
+  });
+
+  it("should reject an overflowing soft_budget instead of silently clearing it", async () => {
+    renderWithSoftBudget(25);
+
+    fireEvent.click(screen.getByText("Settings"));
+    fireEvent.click(screen.getByText("Edit Settings"));
+    (globalThis as any).__TEST_FORM_VALUES = {
+      token: "tok_123",
+      soft_budget: "1e309",
+    };
+
+    fireEvent.click(screen.getByText("Mock Submit"));
+
+    await waitFor(() => expect(toast.error).toHaveBeenCalled());
+    expect(keyUpdateCallMock).not.toHaveBeenCalled();
+  });
+});

ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts

@@ -22,6 +22,7 @@ export interface KeyEditFormValues {
   models?: string[];
   allowed_routes?: string;
   max_budget?: number | string | null;
+  soft_budget?: number | string | null;
   budget_duration?: string | null;
   tpm_limit?: number | string | null;
   tpm_limit_type?: string | null;
@@ -67,6 +68,8 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues =>
   allowed_routes:
     Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0 ? keyData.allowed_routes.join(", ") : "",
   max_budget: keyData.max_budget,
+  soft_budget:
+    (keyData.litellm_budget_table as { soft_budget?: number | null } | null | undefined)?.soft_budget ?? null,
   budget_duration: canonicalBudgetDuration(keyData.budget_duration),
   tpm_limit: keyData.tpm_limit,
   tpm_limit_type: (keyData as { tpm_limit_type?: string | null }).tpm_limit_type ?? null,
@@ -117,6 +120,7 @@ export const keyEditFormSchema = z.object({
   models: z.custom<string[] | undefined>(),
   allowed_routes: z.custom<string | undefined>(),
   max_budget: z.custom<number | string | null | undefined>(),
+  soft_budget: z.custom<number | string | null | undefined>(),
   budget_duration: z.custom<string | null | undefined>(),
   tpm_limit: z.custom<number | string | null | undefined>(),
   tpm_limit_type: z.custom<string | null | undefined>(),
@@ -168,6 +172,7 @@ export const toSubmittedValues = (
   models: values.models,
   allowed_routes: values.allowed_routes,
   max_budget: values.max_budget,
+  soft_budget: values.soft_budget,
   budget_duration: values.budget_duration,
   tpm_limit: values.tpm_limit,
   tpm_limit_type: values.tpm_limit_type,

ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx

@@ -1885,6 +1885,7 @@ describe("KeyEditView", () => {
     key_alias: "asdasdas",
     models: [],
     max_budget: 0,
+    soft_budget: null,
     budget_duration: "30d",
     tpm_limit: 10,
     tpm_limit_type: null,

ui/litellm-dashboard/src/components/templates/key_edit_view.tsx

@@ -32,7 +32,7 @@ import {
   modelSentinelOptions,
   parseAllowedRoutes,
 } from "./keyEditFieldNormalizers";
-import { KeyTypeSelect, labelWithHint } from "./KeyEditViewControls";
+import { KeyBudgetNumberField, KeyTypeSelect, labelWithHint } from "./KeyEditViewControls";
 import {
   AgentsAndGroups,
   KeyEditFormValues,
@@ -417,17 +417,19 @@ export function KeyEditView({
             )}
           </FormField>
 
-          <FormField control={form.control} name="max_budget" label="Max Budget (USD)">
-            {({ ref: _ref, ...field }) => (
-              <NumericalInput
-                {...field}
-                value={field.value ?? ""}
-                step={0.01}
-                style={{ width: "100%" }}
-                placeholder="Enter a numerical value"
-              />
-            )}
-          </FormField>
+          <KeyBudgetNumberField
+            control={form.control}
+            name="max_budget"
+            label="Max Budget (USD)"
+            placeholder="Enter a numerical value"
+          />
+
+          <KeyBudgetNumberField
+            control={form.control}
+            name="soft_budget"
+            label="Soft Budget (USD)"
+            placeholder="Get alerts when spend crosses this value, without blocking requests"
+          />
 
           <FormField control={form.control} name="budget_duration" label="Reset Budget">
             {({ value, onChange, id }) => (

ui/litellm-dashboard/src/components/templates/key_info_view.tsx

@@ -218,6 +218,23 @@ export default function KeyInfoView({
       // Handle max budget empty string
       formValues.max_budget = mapEmptyStringToNull(formValues.max_budget);
 
+      // soft_budget is a budget change server-side (admin-gated); only send it when it changed
+      // so a non-admin edit of unrelated fields isn't blocked by that gate.
+      const previousSoftBudget =
+        (currentKeyData.litellm_budget_table as { soft_budget?: number | null } | null | undefined)?.soft_budget ??
+        null;
+      const nextSoftBudget =
+        formValues.soft_budget === "" || formValues.soft_budget == null ? null : Number(formValues.soft_budget);
+      if (nextSoftBudget !== null && !Number.isFinite(nextSoftBudget)) {
+        toast.error("Soft Budget must be a finite number");
+        return;
+      }
+      if (nextSoftBudget === previousSoftBudget) {
+        delete formValues.soft_budget;
+      } else {
+        formValues.soft_budget = nextSoftBudget;
+      }
+
       // Handle object_permission updates
       if (formValues.vector_stores !== undefined) {
         formValues.object_permission = {

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -8039,7 +8039,7 @@ export interface paths {
          *     - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
          *     - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
          *     - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
-         *     - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
+         *     - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget.
          *     - max_parallel_requests: Optional[int] - Rate limit for parallel requests
          *     - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"}
          *     - tpm_limit: Optional[int] - Tokens per minute limit
@@ -37761,6 +37761,8 @@ export interface components {
             rpm_limit?: number | null;
             /** Rpm Limit Type */
             rpm_limit_type?: ("guaranteed_throughput" | "best_effort_throughput" | "dynamic") | null;
+            /** Soft Budget */
+            soft_budget?: number | null;
             /** Spend */
             spend?: number | null;
             /** Tag Rpm Limit */