fix: /key/block and /key/unblock return 404 (not 401) for non-existent keys#23977
Conversation
… for non-existent keys The block_key() and unblock_key() handlers previously returned a misleading 401 'Authentication Error' when the body 'key' didn't exist in the database, even though authentication (via Authorization header) succeeded correctly. Root cause: After auth passed, the handlers called get_key_object() for cache refresh. This function was designed for auth token lookup and raises ProxyException(code=401) when a token isn't found. Additionally, Prisma's update() silently returns None for non-existent records instead of raising an error, so the code reached get_key_object() without detecting the missing key. Fix: - Add an explicit existence check (find_unique) before the update - Return 404 ProxyException with 'Key not found' if the key doesn't exist - Replace get_key_object() + manual cache update with _delete_cache_key_object() to invalidate the cache (next read will re-fetch from DB) - Reuse the find_unique result for audit logs, eliminating duplicate queries Co-authored-by: yuneng-jiang <[email protected]>
- test_block_key_nonexistent_key_returns_404: verifies block_key returns 404 (not misleading 401) when the key doesn't exist in the DB - test_unblock_key_nonexistent_key_returns_404: same for unblock_key - test_block_key_existing_key_succeeds: verifies block_key succeeds and invalidates cache for existing keys - Update test_unblock_key_supports_both_sk_and_hashed_tokens to reflect the new cache invalidation pattern (_delete_cache_key_object instead of get_key_object + _cache_key_object) Co-authored-by: yuneng-jiang <[email protected]>
These were only used in block_key/unblock_key for cache refresh, which now uses _delete_cache_key_object instead. Co-authored-by: yuneng-jiang <[email protected]>
Co-authored-by: yuneng-jiang <[email protected]>
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
|
Greptile SummaryThis PR fixes a misleading Key changes:
Notable concerns:
Confidence Score: 4/5
|
| Filename | Overview |
|---|---|
| litellm/proxy/management_endpoints/key_management_endpoints.py | Fixes 401-for-missing-key bug in block_key/unblock_key by adding an explicit find_unique existence check; replaces get_key_object+_cache_key_object cache update pattern with _delete_cache_key_object; also fixes a pre-existing bug where unblock_key was logging audit action as "blocked" instead of "unblocked". For proxy-admin callers, a new extra DB round-trip is introduced since _check_key_admin_access already queries the same row for non-proxy-admin callers. |
| tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py | Adds three new focused unit tests for the 404 fix (nonexistent key for block, nonexistent key for unblock, existing key success path); updates existing helpers (_setup_block_unblock_mocks and test_unblock_key_supports_both_sk_and_hashed_tokens) to replace get_key_object+_cache_key_object mocks with _delete_cache_key_object mock; cache invalidation is mocked but not asserted in the success-path test. |
Sequence Diagram
sequenceDiagram
participant C as Client
participant E as block_key / unblock_key
participant A as _check_key_admin_access
participant DB as Prisma DB
participant Cache as user_api_key_cache
C->>E: POST /key/block (Authorization: master key)
E->>A: check admin access (proxy_admin skips, others query DB)
alt Non-proxy-admin: key not found
A->>DB: find_unique(hashed_token)
DB-->>A: None
A-->>E: HTTPException(404)
E-->>C: 404 (HTTPException format)
end
Note over A,E: Proxy-admin returns immediately (no DB check)
E->>DB: find_unique(hashed_token) [new existence check]
alt Key not found
DB-->>E: None
E-->>C: 404 ProxyException "Key not found"
end
DB-->>E: existing_record
E->>DB: update(blocked=True)
DB-->>E: updated record
E->>Cache: _delete_cache_key_object(hashed_token)
Cache-->>E: cache invalidated
E-->>C: 200 updated record
Comments Outside Diff (1)
-
litellm/proxy/management_endpoints/key_management_endpoints.py, line 4977-4988 (link)Unhandled
Nonereturn fromupdateon concurrent deletionBetween the
find_uniquecheck (line 4947) and theupdatecall here, the key could be deleted by a concurrent request. As the PR description notes, Prisma'supdatereturnsNonewhen no row is matched, sorecordwould beNoneand the endpoint would returnHTTP 200with anullbody rather than surfacing an error to the caller.Adding a post-update guard (
if record is None: raise ProxyException(...)) would prevent this silent success-with-null response. The same race applies tounblock_keyat line 5086.
Last reviewed commit: "merge: resolve confl..."
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if token.startswith("sk-"): | ||
| return test_hashed_token | ||
| return token | ||
|
|
There was a problem hiding this comment.
test_block_key_existing_key_succeeds (and the updated test_unblock_key_supports_both_sk_and_hashed_tokens) mock _delete_cache_key_object with a no-op but never assert it was actually invoked. If the code path accidentally skips cache invalidation, the tests would still pass undetected.
Consider using an AsyncMock and asserting it was called:
mock_delete_cache = AsyncMock()
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
mock_delete_cache,
)
# ... after calling block_key / unblock_key ...
mock_delete_cache.assert_called_once()
yuneng-jiang
changed the base branch from
litellm_yj_march_17_2026
to
litellm_yj_march_18_2026
Incorporate new _check_key_admin_access() calls from the base branch into block_key/unblock_key alongside our existence-check fix. Update test mocks: replace references to removed get_key_object and _cache_key_object with _delete_cache_key_object in both the shared _setup_block_unblock_mocks helper and individual test functions. Co-authored-by: yuneng-jiang <[email protected]>
| # Check if the key exists before trying to block it | ||
| existing_record = await prisma_client.db.litellm_verificationtoken.find_unique( | ||
| where={"token": hashed_token} | ||
| ) | ||
| if existing_record is None: | ||
| raise ProxyException( | ||
| message=f"Key not found. Passed key={data.key}", | ||
| type=ProxyErrorTypes.not_found_error, | ||
| param="key", | ||
| code=status.HTTP_404_NOT_FOUND, | ||
| ) |
There was a problem hiding this comment.
find_unique DB call for non-proxy-admin users
_check_key_admin_access (called just above) already issues an identical find_unique query for any caller who is not a proxy admin (see lines 4852–4858). This means team admins and org admins that pass the access-control check will now hit the database twice with the same query — once in _check_key_admin_access, and again with the new existing_record check here — before reaching the update call.
The new existence guard is only necessary for proxy admins, who return early at line 4848 and skip that lookup entirely. For all other callers the key's existence (and the HTTPException(404)) is already handled inside _check_key_admin_access.
The simplest targeted fix is to gate the new check on the caller being a proxy admin, or refactor _check_key_admin_access to return the fetched key row so it can be reused for the audit log below.
The same redundancy applies to unblock_key at line 5062.
Accept theirs for UI test file conflicts (not related to our changes). Our key_management_endpoints.py merged cleanly. Co-authored-by: yuneng-jiang <[email protected]>
| # Check if the key exists before trying to unblock it | ||
| existing_record = await prisma_client.db.litellm_verificationtoken.find_unique( | ||
| where={"token": hashed_token} | ||
| ) | ||
| if existing_record is None: | ||
| raise ProxyException( | ||
| message=f"Key not found. Passed key={data.key}", | ||
| type=ProxyErrorTypes.not_found_error, | ||
| param="key", | ||
| code=status.HTTP_404_NOT_FOUND, | ||
| ) |
There was a problem hiding this comment.
Non-proxy-admin callers with a missing key hit _check_key_admin_access, which raises HTTPException(404, {"error": "Key not found: {hashed_token}"}) — exposing the hashed token in the error detail.
Proxy-admin callers skip that check and now reach this new guard, which raises ProxyException(404) — exposing the original key value from the request body.
The two paths surface different exception types (HTTPException vs ProxyException), different message formats, and different key representations for the same logical error. Aligning both to the same exception type and message template would make client-side error handling consistent. The same inconsistency exists in block_key at lines 4946–4956.
yuneng-jiang
merged commit 7cbabd9
into
litellm_yj_march_18_2026
The /key/update endpoint's get_data() call raises a 401 when the body `key` field doesn't exist in the DB, because get_data() treats the token as an auth credential. This caused the auth layer to resolve the body key instead of the Authorization header bearer token. Replace prisma_client.get_data() with direct Prisma find_unique() in both _get_and_validate_existing_key() and update_key_fn(), matching the pattern used in the /key/block and /key/unblock fix (PR #23977). Also fix the incorrect "Team not found" error message in update_key_fn. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
3 participants
Relevant issues
Fixes issue where
/key/block,/key/unblockreturn401 Authentication Errorfor a non-existent bodykeyeven when theAuthorizationheader contains a valid master key.Pre-Submission checklist
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewType
🐛 Bug Fix
Changes
Problem
Despite the caller authenticating correctly with a valid master key,
/key/blockand/key/unblockreturned401 Authentication Errorwhen the bodykeydidn't exist in the database.Root Cause
The bug was NOT in the auth layer. Authentication passed correctly via the
Authorization: Bearerheader. The401came from inside the handler functions:block_key()/unblock_key()calledprisma_client.db.litellm_verificationtoken.update()which silently returnsNonefor non-existent records (Prisma doesn't error on missing rows forupdate).get_key_object()for cache refresh — this function was designed for auth token lookup and raisesProxyException(code=401, message="Authentication Error...")when a key isn't found.401with "Authentication Error" message propagated to the client, making it look like an auth failure.Fix
find_unique) before the update operation in bothblock_key()andunblock_key()404 Not Foundwith"Key not found. Passed key=..."when the key doesn't existget_key_object()+_cache_key_object()cache refresh with_delete_cache_key_object()(invalidate cache so next read re-fetches from DB)find_uniqueresult for audit logs, eliminating duplicate queriesAfter Fix
Tests Added
test_block_key_nonexistent_key_returns_404— verifies 404 (not 401) for missing keytest_unblock_key_nonexistent_key_returns_404— same for unblocktest_block_key_existing_key_succeeds— verifies success path + cache invalidationtest_unblock_key_supports_both_sk_and_hashed_tokensfor new cache patternAll 161 tests in
test_key_management_endpoints.pypass.