rename

backend/chainlit/auth/__init__.py

@@ -7,7 +7,7 @@
 from chainlit.logger import logger
 from chainlit.oauth_providers import get_configured_oauth_providers
 
-from .cookie import OAuth2PasswordBearerWithCookie, reconstruct_token_from_cookies
+from .cookie import OAuth2PasswordBearerWithCookie, get_token_from_cookies
 from .jwt import create_jwt, decode_jwt, get_jwt_secret
 
 reuseable_oauth = OAuth2PasswordBearerWithCookie(tokenUrl="/login", auto_error=False)
@@ -84,5 +84,5 @@ async def get_current_user(token: str = Depends(reuseable_oauth)):
     "create_jwt",
     "get_configuration",
     "get_current_user",
-    "reconstruct_token_from_cookies",
+    "get_token_from_cookies",
 ]

backend/chainlit/auth/cookie.py

@@ -46,7 +46,7 @@ def __init__(
 
     async def __call__(self, request: Request) -> Optional[str]:
         # First try to get the token from the cookie
-        token = reconstruct_token_from_cookies(request.cookies)
+        token = get_token_from_cookies(request.cookies)
 
         # If no cookie, try the Authorization header as fallback
         if not token:
@@ -76,7 +76,7 @@ async def __call__(self, request: Request) -> Optional[str]:
         return token
 
 
-def reconstruct_token_from_cookies(request_cookies: dict) -> Optional[str]:
+def get_token_from_cookies(cookies: dict[str, str]) -> Optional[str]:
     """
     Read all chunk cookies and reconstruct the token
     """
@@ -86,26 +86,36 @@ def reconstruct_token_from_cookies(request_cookies: dict) -> Optional[str]:
     i = 0
     while True:
         cookie_key = f"{_auth_cookie_name}_{i}"
-        if cookie_key not in request_cookies:
+        if cookie_key not in cookies:
             break
-        chunk_parts.append(request_cookies[cookie_key])
+        chunk_parts.append(cookies[cookie_key])
         i += 1
 
     joined = "".join(chunk_parts)
 
     return joined if joined != "" else None
 
 
-def set_auth_cookie(response: Response, token: str):
+def set_auth_cookie(response: Response, token: str, request: Optional[Request] = None):
     """
     Helper function to set the authentication cookie with secure parameters
+    and remove any leftover chunks from a previously larger token.
     """
-
     _chunk_size = 3000
-    # Split the token into multiple chunks of up to _chunk_size
     chunks = [token[i : i + _chunk_size] for i in range(0, len(token), _chunk_size)]
 
-    # For each chunk, set a separate cookie: auth_chunk_0, auth_chunk_1, etc.
+    # First, delete any old leftover chunk cookies.
+    # If we have the request, we can see exactly which chunk cookies exist.
+    if request is not None:
+        i = 0
+        while True:
+            old_key = f"{_auth_cookie_name}_{i}"
+            if old_key not in request.cookies:
+                break
+            response.delete_cookie(key=old_key, path="/")
+            i += 1
+
+    # Now set the new chunks
     for i, chunk in enumerate(chunks):
         cookie_key = f"{_auth_cookie_name}_{i}"
         response.set_cookie(
@@ -117,7 +127,6 @@ def set_auth_cookie(response: Response, token: str):
             max_age=config.project.user_session_timeout,
         )
 
-
 def clear_auth_cookie(request: Request, response: Response):
     """
     Helper function to clear the authentication cookie

backend/chainlit/socket.py

@@ -8,7 +8,7 @@
 
 from chainlit.auth import (
     get_current_user,
-    reconstruct_token_from_cookies,
+    get_token_from_cookies,
     require_login,
 )
 from chainlit.chat_context import chat_context
@@ -87,7 +87,7 @@ def load_user_env(user_env):
 def _get_token_from_cookie(environ: WSGIEnvironment) -> Optional[str]:
     if cookie_header := environ.get("HTTP_COOKIE", None):
         cookies = cookie_parser(cookie_header)
-        return reconstruct_token_from_cookies(cookies)
+        return get_token_from_cookies(cookies)
 
     return None