fix(proxy): centralize key mutation authorization (LIT-4072)

[yucheng-berri]

Relevant issues

Supersedes #31469. That PR proved the original VERIA-392 finding on /key/generate was just the visible edge of a broader policy drift across six write paths. The iterative review on it surfaced real edge cases (NaN bypass, CLI session token reading None as unlimited, explicit-empty permissions clearing admin-set capabilities, unchanged prefilled values failing the ceiling, bulk paths skipping the gate) but the resulting 17-commit history was unreviewable.

This PR ships the same security outcome as a single coherent change.

Linear ticket

Resolves LIT-4072

Pre-Submission checklist

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Type

🐛 Bug Fix

Changes

Six write paths each enforced subsets of the key-mutation policy inline at their own call sites: /key/generate, /key/service-account/generate, /key/update, /key/regenerate, /key/bulk_update, /team/key/bulk_update. The drift produced real bypasses

This PR collapses the policy into one helper, authorize_key_mutation, that diffs the request against the existing key and enforces three guards by changed field

Rule 1 — permissions is proxy-admin-only. On create paths, only a non-empty dict trips the gate (preserving the empty {} default as the legitimate non-admin shape). On update / regenerate / bulk paths, any explicit presence in model_fields_set trips the gate, so {} and null cannot clear an admin-set capability such as enable_llm_guard_check

Rule 2 — on update-like paths, any budget-field change (max_budget / spend / budget_limits) requires key-admin authority (proxy admin, team admin of the key's team, or org admin of that team's org). Personal-key-owner and team-member-with-grant fast paths apply to non-budget changes only

Rule 3 — every write path enforces numeric hygiene first (NaN / inf / None / wrong-type rejected with 400 before any comparison) and then the delegation ceiling. Hygiene applies to all callers, including proxy admins, because a NaN at rest disables downstream enforcement (spend > NaN is always False)

Rule 4 — ceiling derived from caller.max_budget; falls back to team budget when the caller is a CLI session token in a team context. A CLI session token with no team and any submitted budget is hard-rejected rather than treated as unlimited delegation authority

Rule 5 — delegation_ceiling=None after rule 4 means no ceiling configured (admin granted unlimited delegation); proxy admins and the UI team-admin sentinel session (UI_SESSION_TOKEN_TEAM_ID) are exempt from ceiling enforcement

Rule 6 — on update-like paths, the ceiling is enforced only on values that actually changed against the existing key row. A team admin with max_budget=$100 whose UI prefills the existing $1M cap while editing an unrelated field passes; raising it to $2M is rejected

Rule 7 — generate_key_helper_fn's permissions: Optional[PermissionsDict] = {} mutable default is replaced with None and normalized at the json.dumps call site so the DB shape is preserved

The helper lives in litellm/proxy/management_helpers/key_mutation_authz.py as a single file with one public entrypoint. The six handlers each call it once, in the position where they already have existing_key_row (None on create) and team_table resolved

Screenshots / Proof of Fix

Setup: admin creates user alice-newpr with role internal_user, max_budget=10, models ["gpt-3.5-turbo"], and a personal key for her

1. /key/generate budget_limits over ceiling (expect 400)
{
  "error": {
    "message": "{'error': \"budget_limits entry max_budget (1000000.0) cannot exceed the caller's own max_budget (10.0).\"}",
    "code": "400"
  }
}

2. /key/generate NaN window (expect 400)
{
  "error": {
    "message": "{'error': 'budget_limits entry max_budget must be a finite number; got nan'}",
    "code": "400"
  }
}

3. /key/generate non-admin permissions self-grant (expect 403)
{
  "error": {
    "message": "{'error': 'Only proxy admins can write `permissions` on a key.'}",
    "code": "403"
  }
}

4. /key/update non-admin owner sets permissions (expect 403)
{
  "error": {
    "message": "{'error': 'Only proxy admins can write `permissions` on a key.'}",
    "code": "403"
  }
}

5. /key/update non-admin owner clears permissions explicitly (expect 403)
{
  "error": {
    "message": "{'error': 'Only proxy admins can write `permissions` on a key.'}",
    "code": "403"
  }
}

CONTROL alice within her ceiling (expect 200)
{"key":"sk-E1ljL0dt9Tm...","budget_limits":[{"budget_duration":"1d","max_budget":5.0,"reset_at":"2026-06-29T00:00:00Z"}]}

CONTROL admin can set anything (expect 200)
{"key":"sk-...","permissions":null,"budget_limits":null}

/key/regenerate is enterprise-gated upstream, so the proxy returns the existing enterprise-license error before the helper runs in this non-premium dev setup; the helper still fires on premium

33 new policy-helper tests in tests/test_litellm/proxy/management_helpers/test_key_mutation_authz.py cover the rule matrix end to end. Mutation-killed: 20 of 33 tests fail when the helper is replaced with a no-op. One pre-existing bulk-update test was encoding the buggy "team member with grant can bulk-edit budget" behavior; its update payload was switched to a non-budget field (tpm_limit) to assert the legitimate semantic

336 tests pass across the new helper module and the existing key-management endpoint tests

---

Note

High Risk
Touches authentication-adjacent key management and closes known authorization bypasses; behavior changes on bulk budget edits and permissions writes for non-admins.

Overview
Centralizes proxy key-write authorization in authorize_key_mutation (key_mutation_authz.py) and wires it into generate, update, bulk update, and regenerate instead of scattered inline checks.

Policy enforced in one place: proxy-admin-only permissions (including blocking non-admins from clearing capabilities via {}/null on updates); budget fields on update paths require key-admin authority while personal-owner / team-member fast paths stay limited to non-budget edits; finite-number validation on budgets; delegation ceiling against the caller (with CLI session-token and UI team-admin carve-outs); ceiling and admin gates only on changed budget values on updates (including temp_budget_increase effective cap).

Supporting changes: PermissionsDict typing; permissions default None with {} preserved at JSON serialization; large dedicated test matrix plus bulk team test adjusted to non-budget field.

^{Reviewed by Cursor Bugbot for commit a2ef846. Bugbot is set up for automated code reviews on this repo. Configure here.}

[codecov]

Codecov Report

❌ Patch coverage is 91.86992% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...y/management_endpoints/key_management_endpoints.py	50.00%	5 Missing ⚠️
...llm/proxy/management_helpers/key_mutation_authz.py	95.41%	5 Missing ⚠️

📢 Thoughts on this report? Let us know!

[yucheng-berri]

Single coherent commit implementing the centralized policy from LIT-4072, plus a follow-up that types the permissions field with PermissionsDict mirroring the ObjectPermissionDict pattern in #31471. Supersedes the now-closed #31469. Live verify ran end-to-end on the attack matrix; 337 tests pass.

@greptileai please review
@veria-ai please review

[veria-ai]

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 2 · PR risk: 0/10

[greptile-apps]

Greptile Summary

This PR centralizes key mutation authorization for key-management write paths. The main changes are:

• Adds authorize_key_mutation for permissions, budget authority, numeric hygiene, and delegation ceiling checks
• Wires the helper into key generation, key update, key regeneration, and bulk key update paths
• Replaces the mutable permissions default in generate_key_helper_fn while preserving the stored DB shape
• Narrows key permissions typing with PermissionsDict
• Adds focused tests for the new helper and updates bulk team-key tests for the new budget-edit semantics

Confidence Score: 4/5

The authorization-sensitive changes are focused and backed by a broad helper test matrix, with no outstanding code issues identified.

The implementation consolidates duplicated policy checks into a single helper and updates the relevant write paths and tests, reducing drift while preserving expected database shapes.

No specific files require follow-up from this review.

T-Rex Logs

What T-Rex did

• Compared the generate-authz baseline before (trex-artifacts/generate-authz-01-before.log) and after (trex-artifacts/generate-authz-02-after.log) artifacts to verify the behavioral shift.
• Examined the central helper behavior in the generate-authz scenario, noting base runs show helper_missing and head runs reveal additional 403/400 outcomes across permission and budget controls.
• Compared the bulk-authz baseline (bulk-authz-01-before.log) and after (bulk-authz-02-after.log) artifacts to verify bulk changes now yield 403 for budget and permission updates, with NaN max_budget treated as 400 and non-budget tpm_limit remaining 200 OK.
• Reviewed the permissions insert tests before and after, confirming the default and explicit payload representations and that both runs exited with exit code 0.

_{Ran code and verified through T-Rex}

_{Reviews (4): Last reviewed commit: "fix: honor existing key team for ui auth..." | Re-trigger Greptile}

[greptile-apps]

Greptile Summary

This PR centralizes key mutation authorization across the key management write paths. The main changes are:

• Adds authorize_key_mutation for key creation, update, regenerate, and bulk update policy checks.
• Enforces permission writes, budget delegation ceilings, and numeric budget hygiene in one helper.
• Updates key management handlers to call the helper before mutation.
• Replaces the mutable permissions default in generate_key_helper_fn.
• Adds tests for the helper policy matrix and updates bulk update endpoint coverage.

Confidence Score: 3/5

The change improves centralized authorization, but one update path still appears to allow explicit budget-cap removal by actors who should not be able to make budget changes.

The modified policy is well covered in many paths, but the null update case needs attention before this is safe to merge.

litellm/proxy/management_helpers/key_mutation_authz.py

T-Rex Logs

What T-Rex did

• T-Rex attempted a focused reproduction script for explicit max_budget=None, but the required key_mutation_authz helper module was not present in the workspace, blocking execution and leaving reproduction artifacts saved for rerun.
• T-Rex ran probes of prohibited /key/generate cases and observed that the head still returned 200 OK for these cases, indicating the behavior persisted.
• T-Rex executed the update_key_policy_probe against /key/update and /key/regenerate and found 500 Internal Server Error due to a missing litellm.proxy.management_helpers.key_mutation_authz module in both base/head, showing the helper is absent in this head.
• T-Rex ran bulk-key-policy tests and found that base allowed and changed max_budget values, while head still allowed non-budget settings with failed budget updates and max_budget remaining at 10.0, with the failure path tracing through the authorization checks.
• T-Rex compared permissions-default tests and observed that head returns signature default None and omits persisted permissions, while explicit JSON remains unchanged in both calls, indicating the mutated permissions did not alter those fields.

_{Ran code and verified through T-Rex}

Comments Outside Diff (2)

1. General comment

   Create-key policy helper is not invoked on /key/generate, so non-admin self-grants and budget delegation bypasses remain allowed

   • Bug
     • On head, the executed /key/generate create path still returns 200 OK for requests the PR contract says must be rejected: a non-admin with max_budget=100 can submit budget_limits=[{"budget_duration":"1d","max_budget":500}], can submit budget_limits containing NaN, can submit non-empty permissions={"get_spend_routes":true}, and a CLI session token with no team can submit max_budget=5. These responses contradict the claimed authorization policy changes.
   • Cause
     • The new authorize_key_mutation(...) call exists in key_management_endpoints.py around lines 673-687 inside generate_key_helper_fn, but the actual /key/generate endpoint function generate_key_fn proceeds through validation, key_generation_check, and _common_key_generation_helper around lines 1370-1476 without calling authorize_key_mutation. As a result, helper checks such as _check_permissions_field and _check_delegation_ceiling in key_mutation_authz.py are never reached for this create endpoint in the exercised path.
   • Fix
     • Call authorize_key_mutation(data=data, existing_key_row=None, user_api_key_dict=user_api_key_dict, team_table=team_table, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, route_label="/key/generate") from generate_key_fn after resolving team_table and before _common_key_generation_helper, or refactor /key/generate to consistently go through generate_key_helper_fn where the authorization call already exists. Add endpoint-level regression coverage for these exact create scenarios.

   _{Ran code and verified through T-Rex}

2. General comment

   PR-described key mutation authorization helper is absent from head

   • Bug
     • The validation objective requires /key/update and /key/regenerate to enforce permissions, budget admin gating, numeric hygiene, and delegation ceilings through litellm/proxy/management_helpers/key_mutation_authz.py. Runtime validation against the requested head failed immediately because that module is not present/importable. A direct tree/callsite check also found no authorize_key_mutation reference in key_management_endpoints.py, so the changed authorization contract cannot execute for either route.
   • Cause
     • The head checkout does not include the claimed key_mutation_authz.py helper or endpoint integration callsite, despite the task listing it as a changed file and relying on it for the update/regenerate policy.
   • Fix
     • Add litellm/proxy/management_helpers/key_mutation_authz.py to the commit and wire /key/update and /key/regenerate through authorize_key_mutation before mutation persistence, then rerun the before/after endpoint/helper scenarios.

   _{Ran code and verified through T-Rex}

_{Reviews (2): Last reviewed commit: "refactor(proxy): type permissions field ..." | Re-trigger Greptile}

[yucheng-berri]

@veria-ai you were right on both counts. Pushed a2ef846.

Part 1 (max_budget=null bypass): is_budget_change used to gate only when data.max_budget is not None and != existing. Now keyed off model_fields_set so an explicit null clear by a non-admin owner trips the gate; an unchanged-value resubmit still passes the personal-key fast path.

Part 2 (temp_budget_increase / temp_budget_expiry): both added to is_budget_change so a non-admin can't sneak a temp bump. temp_budget_increase also goes through math.isfinite (same NaN concern as max_budget) and the delegation ceiling is applied to the effective cap (existing + temp_increase) so a team admin with max_budget=$100 can't bump a key from $50 to $1,000,050.

342 tests pass; 4 of 5 new tests mutation-killed against the prior commit (the fifth is a counter-test that pins unchanged-resubmit semantics).

@greptileai please re-review HEAD a2ef846
@veria-ai please re-review HEAD a2ef846

[yucheng-berri]

@veria-ai your review at 02:04 UTC was on commit b270366, which predated my fix at 02:14 UTC. The "Budget mutation bypass" you flagged at key_mutation_authz.py:189 is closed in HEAD a2ef846 with all three concerns addressed (explicit max_budget=null gate via model_fields_set, temp_budget_increase / temp_budget_expiry added to is_budget_change, delegation ceiling applied to the effective max_budget). Please re-review the current HEAD.

[yucheng-berri]

bugbot run

[yucheng-berri]

@veria-ai you were right again. Pushed c8583e3. The ceiling check was using data.max_budget is not None to detect changes, so an explicit max_budget: null from a team / org admin never reached the comparison and silently removed the existing cap.

Now derives max_budget_changed from model_fields_set so the null counts as a change. When the new value is null and an existing finite cap was set, the resulting effective max is unbounded and rejects against any finite delegation ceiling. Proxy admin remains exempt. Counter-tests pin the no-op-clear and admin-clear semantics.

345 tests pass; mutation-killed against the prior commit.

@veria-ai please re-review HEAD c8583e3
@greptileai please re-review HEAD c8583e3

[cursor]

Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.

Bugbot Autofix resolved 1 of the 2 issues found in the latest run.

• ✅ Fixed: UI session misses key team
  • The UI team-admin exemption now resolves the target team from the request body or the existing key row, covering regenerate payloads that omit team_id.

_{You can send follow-ups to the cloud agent here.}

^{Reviewed by Cursor Bugbot for commit a2ef846. Configure here.}

[cursor]

Null max_budget skips ceiling

High Severity

_check_delegation_ceiling treats an explicit max_budget clear as unchanged because it keys off data.max_budget is not None, while _check_budget_admin_authority correctly treats the same payload as a budget change via model_fields_set. A bounded team or org admin who passes the admin gate can remove a key’s cap without any delegation check, effectively granting unlimited spend authority beyond their own max_budget.

Additional Locations (1)

• litellm/proxy/management_helpers/key_mutation_authz.py#L195-L198

 

^{Reviewed by Cursor Bugbot for commit a2ef846. Configure here.}

[cursor]

Bugbot Autofix determined this is a false positive.

Current code already treats explicit max_budget clears as changed and rejects finite-ceiling callers who would remove an existing cap.

_{You can send follow-ups to the cloud agent here.}

[cursor]

UI session misses key team

Medium Severity

The UI team-admin session carve-out in _check_delegation_ceiling only looks at data.team_id, not the target key’s team when the body omits it. Handlers such as /key/regenerate still load team_table from existing_key_row.team_id, but the sentinel session loses Rule 5 exemption and can be blocked by the delegation ceiling on legitimate team-key edits.

Additional Locations (1)

• litellm/proxy/management_helpers/key_mutation_authz.py#L63-L71

 

^{Reviewed by Cursor Bugbot for commit a2ef846. Configure here.}

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ yucheng-berri
❌ cursoragent
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[yucheng-berri]

Closing to revert to the focused scope. The centralized helper grew the same way #31469 did — every bot finding was real but each follow-up pushed the PR further from the original VERIA-392 ticket. Resetting #31469 to its first commit (just the /key/generate gates) and filing the deferred findings as separate Linear tickets so each lands as its own focused PR.