From b7c5073d28ab6f92d322d31775aa6005c372fc39 Mon Sep 17 00:00:00 2001
From: Krish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 29 Jul 2025 16:28:25 -0700
Subject: [PATCH] Custom Auth - bubble up custom exceptions  (#13093)

* fix(enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py): bubble up exception if type is ProxyException

* docs(custom_auth.md): doc on bubbling up custom exceptions
---
 docs/my-website/docs/proxy/custom_auth.md     | 105 ++++++++++++++++++
 .../proxy/auth/user_api_key_auth.py           |   4 +-
 litellm/proxy/custom_auth_auto.py             |  11 +-
 3 files changed, 118 insertions(+), 2 deletions(-)

diff --git a/docs/my-website/docs/proxy/custom_auth.md b/docs/my-website/docs/proxy/custom_auth.md
index 646a68b9d3f..3787f9bdd7c 100644
--- a/docs/my-website/docs/proxy/custom_auth.md
+++ b/docs/my-website/docs/proxy/custom_auth.md
@@ -60,9 +60,114 @@ Supported from v1.72.2+
 [Get free 7-day trial key](https://www.litellm.ai/enterprise#trial)
 :::
 
+### Usage
+
+1. Setup custom auth file
+
+```python
+"""
+Example custom auth function.
+
+This will allow all keys starting with "my-custom-key" to pass through.
+"""
+from typing import Union
+
+from fastapi import Request
+
+from litellm.proxy._types import UserAPIKeyAuth
+
+
+async def user_api_key_auth(
+    request: Request, api_key: str
+) -> Union[UserAPIKeyAuth, str]:
+    try:
+        if api_key.startswith("my-custom-key"):
+            return "sk-P1zJMdsqCPNN54alZd_ETw"
+        else:
+            raise Exception("Invalid API key")
+    except Exception:
+        raise Exception("Invalid API key")
+
+```
+
+2. Setup config.yaml
+
+Key change set `mode: auto`. This will check both litellm api key auth + custom auth.
+
 ```yaml
+model_list: 
+  - model_name: "openai-model"
+    litellm_params: 
+      model: "gpt-3.5-turbo"
+      api_key: os.environ/OPENAI_API_KEY
+
 general_settings:
   custom_auth: custom_auth_auto.user_api_key_auth
   custom_auth_settings:
     mode: "auto" # can be 'on', 'off', 'auto' - 'auto' checks both litellm api key auth + custom auth
+```
+
+Flow:
+1. Checks custom auth first
+2. If custom auth fails, checks litellm api key auth
+3. If both fail, returns 401
+
+
+3. Test it! 
+
+```bash
+curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
+-H 'Content-Type: application/json' \
+-H 'Authorization: Bearer sk-P1zJMdsqCPNN54alZd_ETw' \
+-d '{
+    "model": "openai-model",
+    "messages": [
+          {
+            "role": "user",
+            "content": "Hey! My name is John"
+          }
+        ]
+}'
+```
+
+
+
+
+#### Bubble up custom exceptions
+
+If you want to bubble up custom exceptions, you can do so by raising a `ProxyException`.
+
+```python
+"""
+Example custom auth function.
+
+This will allow all keys starting with "my-custom-key" to pass through.
+"""
+
+from typing import Union
+
+from fastapi import Request
+
+from litellm.proxy._types import UserAPIKeyAuth, ProxyException
+
+
+async def user_api_key_auth(
+    request: Request, api_key: str
+) -> Union[UserAPIKeyAuth, str]:
+    try:
+        if api_key.startswith("my-custom-key"):
+            return "sk-P1zJMdsqCPNN54alZd_ETw"
+        if api_key == "invalid-api-key":
+            # raise a custom exception back to the client
+            raise ProxyException(
+                message="Invalid API key",
+                type="invalid_request_error",
+                param="api_key",
+                code=401,
+            )
+        else:
+            raise Exception("Invalid API key")
+    except Exception:
+        raise Exception("Invalid API key")
+
 ```
\ No newline at end of file
diff --git a/enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py b/enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py
index 35b4c2a1f3b..dc9fdeb78e2 100644
--- a/enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py
+++ b/enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py
@@ -3,7 +3,7 @@ from typing import Any, Optional
 from fastapi import Request
 
 from litellm._logging import verbose_proxy_logger
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 
 
 async def enterprise_custom_auth(
@@ -24,6 +24,8 @@ async def enterprise_custom_auth(
     elif custom_auth_settings["mode"] == "auto":
         try:
             return await user_custom_auth(request, api_key)
+        except ProxyException as e:
+            raise e
         except Exception as e:
             verbose_proxy_logger.debug(
                 f"Error in custom auth, checking litellm auth: {e}"
diff --git a/litellm/proxy/custom_auth_auto.py b/litellm/proxy/custom_auth_auto.py
index 47889e61e57..c8991520898 100644
--- a/litellm/proxy/custom_auth_auto.py
+++ b/litellm/proxy/custom_auth_auto.py
@@ -3,11 +3,12 @@ Example custom auth function.
 
 This will allow all keys starting with "my-custom-key" to pass through.
 """
+
 from typing import Union
 
 from fastapi import Request
 
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 
 
 async def user_api_key_auth(
@@ -16,6 +17,14 @@ async def user_api_key_auth(
     try:
         if api_key.startswith("my-custom-key"):
             return "sk-P1zJMdsqCPNN54alZd_ETw"
+        if api_key == "invalid-api-key":
+            # raise a custom exception back to the client
+            raise ProxyException(
+                message="Invalid API key",
+                type="invalid_request_error",
+                param="api_key",
+                code=401,
+            )
         else:
             raise Exception("Invalid API key")
     except Exception: