[Fix] Key Update Endpoint Returns 401 Instead of 404 for Nonexistent Keys by yuneng-jiang · Pull Request #24063 · BerriAI/litellm

yuneng-jiang · 2026-03-18T23:53:47Z

Summary

Failure Path (Before Fix)

/key/update returns 401 for a nonexistent body key even when the Authorization header contains a valid master key. This happens because prisma_client.get_data() treats the token as an auth credential and raises a 401 when it doesn't exist in the DB.

Fix

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 fixes an incorrect "Team not found" error message in update_key_fn.

Testing

curl -s -X POST "http://localhost:4000/key/update" \
  -H "Authorization: Bearer $MASTER_KEY" \
  -H "Content-Type: application/json" \
  -d '{"key":"sk-does-not-exist"}'
# Before: 401 "invalid user key"
# After:  404 "Key not found"

Type

🐛 Bug Fix

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) <noreply@anthropic.com>

vercel · 2026-03-18T23:53:52Z

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Mar 19, 2026 0:05am

greptile-apps · 2026-03-18T23:57:46Z

Greptile Summary

This PR fixes a bug where POST /key/update incorrectly returns HTTP 401 ("invalid user key") instead of HTTP 404 ("Key not found") when the key field in the request body does not exist in the database, even when the caller is authenticated with a valid master key via the Authorization header.

Root cause: prisma_client.get_data() treated the body key as an auth credential and raised a 401 when the token was absent from the database. The fix replaces this with a direct Prisma find_unique() call (consistent with the pattern used in PR #23977 for /key/block and /key/unblock), and surfaces a ProxyException with HTTP 404 through the existing helper _get_and_validate_existing_key().

Key changes:

_get_and_validate_existing_key(): replaced prisma_client.get_data() with _hash_token_if_needed() + direct find_unique(), and swapped HTTPException(404) for ProxyException(404) so it propagates correctly through update_key_fn()'s try/except block (which re-raises ProxyException as-is, but wraps HTTPException into an auth_error)
update_key_fn(): removed the hand-rolled key-lookup block (which had a copy-paste error — "Team not found" message) and delegated to _get_and_validate_existing_key()
Added test_update_key_nonexistent_key_returns_404 to lock in the fix as a mock-only regression test, consistent with the test pattern used in fix: /key/block and /key/unblock return 404 (not 401) for non-existent keys #23977

Confidence Score: 5/5

Safe to merge — the fix is narrowly scoped, correctly propagates a 404 ProxyException through the existing exception handler, and is covered by a new mock-based regression test.
The change replaces two small, well-understood call sites (get_data() → find_unique()) and removes a copy-paste bug ("Team not found" message). Exception flow through the existing try/except in update_key_fn is unambiguously correct: ProxyException is re-raised as-is, preserving the 404 code. The new test mirrors the established block/unblock test pattern and uses no real network calls.
No files require special attention.

Important Files Changed

Filename	Overview
litellm/proxy/management_endpoints/key_management_endpoints.py	Root-cause fix: replaces `get_data()` with `find_unique()` + `_hash_token_if_needed()` in `_get_and_validate_existing_key()`, and consolidates the duplicate key-lookup in `update_key_fn()` to delegate to that helper. Also fixes the stale "Team not found" error message. Exception propagation through the existing `try/except` is correct (`ProxyException` is re-raised as-is).
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py	Adds `test_update_key_nonexistent_key_returns_404`: a pure-mock unit test (no real network calls) consistent with the project's test pattern. Correctly asserts `code == "404"`, message contains "not found", and "Authentication Error" is absent. The assertion matches `ProxyException`'s `self.code = str(code)` semantics.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Client: POST /key/update] --> B{DB connected?}
    B -- No --> C[HTTPException 500]
    B -- Yes --> D[Hash token via SHA-256]
    D --> E[Query DB via find_unique]
    E --> F{Record found?}
    F -- No --> G[Raise ProxyException 404\nNot Found]
    F -- Yes --> H[Validate update data]
    H --> I[Prepare update fields]
    I --> J[Update DB record]
    J --> L[Invalidate cache entry]
    L --> M[Return success response]
    G --> N{Caught in except block}
    N -- ProxyException --> O[Re-raised as-is\nHTTP 404 to client]
    N -- HTTPException --> P[Wrapped as auth_error\nProxyException]

_{Last reviewed commit: "fix: address review ..."}

greptile-apps · 2026-03-18T23:57:53Z

+    from litellm.proxy.proxy_server import hash_token
+
+    if token.startswith("sk-"):
+        hashed_token = hash_token(token=token)
+    else:
+        hashed_token = token
+
+    existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique(
+        where={"token": hashed_token}
    )

    if existing_key_row is None:
-        raise HTTPException(
-            status_code=404,
-            detail={"error": f"Key not found: {token}"},
+        raise ProxyException(
+            message=f"Key not found. Passed key={token}",
+            type=ProxyErrorTypes.not_found_error,
+            param="key",
+            code=status.HTTP_404_NOT_FOUND,
        )

    return existing_key_row


No automated test for the fixed code path

PR #23977 (the block/unblock fix this PR is modelled on) added test_block_key_nonexistent_key_returns_404 and test_unblock_key_nonexistent_key_returns_404 to prove the fix works. This PR has no equivalent test for /key/update.

Without a mock-based unit test that verifies update_key_fn() raises a ProxyException with HTTP 404 (and not 401) when find_unique returns None, there is no automated regression guard for the reported bug. If the project's custom rule requires evidence of resolution, this should be addressed before merge.

A test similar to the block/unblock ones already present in tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py would be sufficient.

greptile-apps · 2026-03-18T23:57:54Z

+        from litellm.proxy.proxy_server import hash_token
+
+        if data.key.startswith("sk-"):
+            hashed_token = hash_token(token=data.key)
+        else:
+            hashed_token = data.key
+
+        existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique(
+            where={"token": hashed_token}
        )

        if existing_key_row is None:
-            raise HTTPException(
-                status_code=404,
-                detail={"error": f"Team not found, passed team_id={data.team_id}"},
+            raise ProxyException(
+                message=f"Key not found. Passed key={data.key}",
+                type=ProxyErrorTypes.not_found_error,
+                param="key",
+                code=status.HTTP_404_NOT_FOUND,
            )


Duplicated key-lookup logic — use _get_and_validate_existing_key()

update_key_fn() contains a hand-rolled copy of exactly the same token-hashing + find_unique + 404-raise pattern that _get_and_validate_existing_key() was just updated to provide. _process_single_key_update() (used by the bulk-update path) already delegates to that helper correctly.

Since _get_and_validate_existing_key() accepts (token, prisma_client), this block can be replaced by a single call to that helper. Having two divergent copies means any future change (e.g. to error message, hash logic, or audit logging) must be made in both places, risking drift.

greptile-apps · 2026-03-18T23:57:55Z

+    from litellm.proxy.proxy_server import hash_token
+
+    if token.startswith("sk-"):
+        hashed_token = hash_token(token=token)
+    else:
+        hashed_token = token


Inline import and manual hashing — use existing helper

The helper _hash_token_if_needed (imported at line 79 from litellm.proxy.utils) already encapsulates the same conditional-hashing guard used here. Replacing this block with a call to that helper avoids the deferred from litellm.proxy.proxy_server import hash_token import and removes duplicated logic.

The same pattern also appears in update_key_fn() at lines 2122–2127 and should be consolidated as well.

…add test - Deduplicate: update_key_fn now delegates to _get_and_validate_existing_key() instead of inlining its own copy of the lookup logic - Use _hash_token_if_needed (already imported at module level) instead of inline `from proxy_server import hash_token` + manual conditional - Fix stale docstring: _get_and_validate_existing_key raises ProxyException, not HTTPException - Add unit test: test_update_key_nonexistent_key_returns_404 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

- Replace prisma_client.get_data() with direct find_unique() in _get_and_validate_existing_key() and update_key_fn() - Fix incorrect error message from 'Team not found' to 'Key not found' Fixes BerriAI#24063

vercel Bot deployed to Preview March 18, 2026 23:55 View deployment

greptile-apps Bot reviewed Mar 18, 2026

View reviewed changes

vercel Bot deployed to Preview March 19, 2026 00:05 View deployment

Jah-yee mentioned this pull request Mar 19, 2026

fix(proxy): use find_unique directly to avoid 401 on nonexistent key #24064

Open

yuneng-jiang merged commit 71cdd91 into litellm_yj_march_18_2026 Mar 19, 2026
55 of 65 checks passed

ishaan-berri deleted the litellm_fix_key_update_404 branch March 26, 2026 22:29

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Fix] Key Update Endpoint Returns 401 Instead of 404 for Nonexistent Keys#24063

[Fix] Key Update Endpoint Returns 401 Instead of 404 for Nonexistent Keys#24063
yuneng-jiang merged 2 commits intolitellm_yj_march_18_2026from
litellm_fix_key_update_404

yuneng-jiang commented Mar 18, 2026

Uh oh!

vercel Bot commented Mar 18, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Mar 18, 2026 •

edited

Loading

Important Files Changed

Uh oh!

greptile-apps Bot Mar 18, 2026

Uh oh!

greptile-apps Bot Mar 18, 2026

Uh oh!

greptile-apps Bot Mar 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

yuneng-jiang commented Mar 18, 2026

Summary

Failure Path (Before Fix)

Fix

Testing

Type

Uh oh!

vercel Bot commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

greptile-apps Bot commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

greptile-apps Bot Mar 18, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Mar 18, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Mar 18, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

vercel Bot commented Mar 18, 2026 •

edited

Loading

greptile-apps Bot commented Mar 18, 2026 •

edited

Loading