From 9852462d1f44b5f45f80e98103d3593c5f0a8846 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 2 Feb 2024 18:40:18 -0800
Subject: [PATCH] (fix) proxy - clear cache after /key/delete

---
 litellm/proxy/proxy_server.py | 33 +++++++++++++++++++++++++++------
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 7498c45ac..f8e501cb6 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -527,7 +527,7 @@ async def user_api_key_auth(
                     )
 
             # Token passed all checks
-            api_key = valid_token.token
+            api_key = hash_token(valid_token.token)
 
             # Add hashed token to cache
             user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60)
@@ -1445,11 +1445,13 @@ async def generate_key_helper_fn(
             saved_token["expires"], datetime
         ):
             saved_token["expires"] = saved_token["expires"].isoformat()
-        user_api_key_cache.set_cache(
-            key=key_data["token"],
-            value=LiteLLM_VerificationToken(**saved_token),  # type: ignore
-            ttl=60,
-        )
+        if key_data["token"] is not None and isinstance(key_data["token"], str):
+            hashed_token = hash_token(key_data["token"])
+            user_api_key_cache.set_cache(
+                key=hashed_token,
+                value=LiteLLM_VerificationToken(**saved_token),  # type: ignore
+                ttl=60,
+            )
         if prisma_client is not None:
             ## CREATE USER (If necessary)
             verbose_proxy_logger.debug(f"prisma_client: Creating User={user_data}")
@@ -2665,13 +2667,32 @@ async def delete_key_fn(data: DeleteKeyRequest):
         HTTPException: If an error occurs during key deletion.
     """
     try:
+        global user_api_key_cache
         keys = data.keys
+        if len(keys) == 0:
+            raise ProxyException(
+                message=f"No keys provided, passed in: keys={keys}",
+                type="auth_error",
+                param="keys",
+                code=status.HTTP_400_BAD_REQUEST,
+            )
 
         result = await delete_verification_token(tokens=keys)
         verbose_proxy_logger.debug("/key/delete - deleted_keys=", result)
 
         number_deleted_keys = len(result["deleted_keys"])
         assert len(keys) == number_deleted_keys
+
+        for key in keys:
+            user_api_key_cache.delete_cache(key)
+            # remove hash token from cache
+            hashed_token = hash_token(key)
+            user_api_key_cache.delete_cache(hashed_token)
+
+        verbose_proxy_logger.debug(
+            f"/keys/delete - cache after delete: {user_api_key_cache.in_memory_cache.cache_dict}"
+        )
+
         return {"deleted_keys": keys}
     except Exception as e:
         if isinstance(e, HTTPException):