fix

Signed-off-by: San Nguyen <vinhsannguyen91@gmail.com>

backend/chainlit/auth/cookie.py

@@ -25,6 +25,8 @@
 _cookie_secure = _cookie_samesite == "none"
 
 _state_cookie_lifetime = 3 * 60  # 3m
+# TODO: Prefix with __Host- to use in conjunction with the partitioned cookie as described here
+# [CHIPS](https://github.com/privacycg/CHIPS/blob/main/README.md)
 _auth_cookie_name = "access_token"
 _state_cookie_name = "oauth_state"
 
@@ -89,6 +91,14 @@ def set_auth_cookie(response: Response, token: str):
         samesite=_cookie_samesite,
         max_age=config.project.user_session_timeout,
     )
+    if _cookie_secure:
+        # Set partitioned cookie to avoid cookie rejection errors
+        # https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie#partitioned
+        # TODO: Use FastAPI native methods once it supports partitioned cookie
+        # https://github.com/fastapi/fastapi/discussions/11285
+        response.headers["Set-Cookie"] = (
+            response.headers["Set-Cookie"] + "; Partitioned"
+        )
 
 
 def clear_auth_cookie(response: Response):
@@ -107,6 +117,14 @@ def set_oauth_state_cookie(response: Response, token: str):
         secure=_cookie_secure,
         max_age=_state_cookie_lifetime,
     )
+    if _cookie_secure:
+        # Set partitioned cookie to avoid cookie rejection errors
+        # https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie#partitioned
+        # TODO: Use FastAPI native methods once it supports partitioned cookie
+        # https://github.com/fastapi/fastapi/discussions/11285
+        response.headers["Set-Cookie"] = (
+            response.headers["Set-Cookie"] + "; Partitioned"
+        )
 
 
 def validate_oauth_state_cookie(request: Request, state: str):

backend/chainlit/cli/__init__.py

@@ -200,6 +200,7 @@ def chainlit_run(
         os.environ["OPENAI_API_KEY"] = "sk-FAKE-OPENAI-API-KEY"
         # This is required for authentication tests
         os.environ["CHAINLIT_AUTH_SECRET"] = "SUPER_SECRET"  # nosec B105
+        os.environ["CHAINLIT_COOKIE_SAMESITE"] = "none"
     else:
         trace_event("chainlit run")
 

backend/chainlit/element.py

@@ -95,7 +95,7 @@ class Element:
     for_id: Optional[str] = None
     # The language, if relevant
     language: Optional[str] = None
-    # Mime type, infered based on content if not provided
+    # Mime type, inferred based on content if not provided
     mime: Optional[str] = None
 
     def __post_init__(self) -> None:

backend/chainlit/socket.py

@@ -7,6 +7,7 @@
 from typing_extensions import TypeAlias
 
 from chainlit.auth import get_current_user, require_login
+from chainlit.auth.cookie import _auth_cookie_name
 from chainlit.chat_context import chat_context
 from chainlit.config import config
 from chainlit.context import init_ws_context
@@ -83,7 +84,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 cookies.get("access_token", None)
+        return cookies.get(_auth_cookie_name, None)
 
     return None