chore: import upstream snapshot with attribution

backend/tests/auth/test_cookie.py

@@ -0,0 +1,165 @@
+import importlib
+
+import pytest
+from fastapi import FastAPI, Form
+from fastapi.testclient import TestClient
+from starlette.requests import Request
+from starlette.responses import Response
+
+import chainlit.auth.cookie as cookie_module
+from chainlit.auth import (
+    clear_auth_cookie,
+    get_token_from_cookies,
+    set_auth_cookie,
+)
+
+
+@pytest.fixture
+def test_app():
+    app = FastAPI()
+
+    @app.post("/set-cookie")
+    async def set_cookie_endpoint(request: Request, token: str = Form()):
+        response = Response()
+        set_auth_cookie(request, response, token)
+        return response
+
+    @app.get("/get-token")
+    async def get_token_endpoint(request: Request):
+        token = get_token_from_cookies(request.cookies)
+        return {"token": token}
+
+    @app.delete("/clear-cookie")
+    async def clear_cookie_endpoint(request: Request):
+        response = Response()
+        clear_auth_cookie(request, response)
+        return response
+
+    return app
+
+
+@pytest.fixture
+def client(test_app):
+    return TestClient(test_app)
+
+
+def test_short_token(client):
+    """Test with a <3000 shorter token."""
+
+    # Set a short token
+    short_token = "x" * 1000
+    set_response = client.post("/set-cookie", data={"token": short_token})
+    assert set_response.status_code == 200
+
+    # Verify cookies were set
+    cookies = set_response.cookies
+    assert cookies, "No cookies set"
+    assert "access_token" in cookies, f"No chunking for short cookies: {cookies}"
+
+    # Read back the token using client's cookie jar
+    get_response = client.get("/get-token")
+    assert get_response.status_code == 200
+    assert get_response.json()["token"] == short_token
+
+
+def test_set_and_read_4kb_token(client):
+    """Test full cookie lifecycle using actual client cookie handling."""
+
+    # Set a 4KB token
+    token_4kb = "x" * 4000
+    set_response = client.post("/set-cookie", data={"token": token_4kb})
+    assert set_response.status_code == 200
+
+    # Verify cookies were set
+    cookies = set_response.cookies
+    assert f"{cookies.keys()} should contain chunked cookies", any(
+        key.startswith("access_token_") for key in cookies.keys()
+    )
+
+    # Read back the token using client's cookie jar
+    get_response = client.get("/get-token")
+    assert get_response.status_code == 200
+
+    response_token = get_response.json()["token"]
+    assert len(response_token) == len(token_4kb)
+    assert response_token == token_4kb
+
+
+def test_overwrite_shorter_token_chunked(client):
+    """Test cookie chunk cleanup when replacing a large token with a smaller one."""
+    # Set initial long token
+    long_token = "LONG" * 2000  # 8000 characters
+    client.post("/set-cookie", data={"token": long_token})
+
+    # Verify initial chunks exist
+    first_cookies = client.cookies
+    assert len([k for k in first_cookies if k.startswith("access_token_")]) > 1
+
+    # Set shorter token (should clear previous chunks)
+    short_token = "SHORT" * 1000  # 4000 characters
+    client.post("/set-cookie", data={"token": short_token})
+
+    # Verify new cookie state
+    final_response = client.get("/get-token")
+    assert final_response.json()["token"] == short_token
+
+    # Verify only two chunks remain
+    final_cookies = client.cookies
+    chunk_cookies = [k for k in final_cookies if k.startswith("access_token_")]
+    assert len(chunk_cookies) == 2, f"Found {len(chunk_cookies)} residual cookies"
+
+
+def test_overwrite_shorter_token_unchunked(client):
+    """Test cookie chunk cleanup when replacing a large token with a smaller one."""
+    # Set initial long token
+    long_token = "LONG" * 1000  # 4000 characters
+    client.post("/set-cookie", data={"token": long_token})
+
+    # Verify initial chunks exist
+    first_cookies = client.cookies
+    assert len([k for k in first_cookies if k.startswith("access_token_")]) > 1
+
+    # Set shorter token (should clear previous chunks)
+    short_token = "SHORT"
+    client.post("/set-cookie", data={"token": short_token})
+
+    # Verify new cookie state
+    final_response = client.get("/get-token")
+    assert final_response.json()["token"] == short_token
+
+    # Verify no chunks remain
+    final_cookies = client.cookies
+    chunk_cookies = [k for k in final_cookies if k.startswith("access_token_")]
+    assert len(chunk_cookies) == 0, f"Found {len(chunk_cookies)} residual cookies"
+
+
+def test_state_cookie_lifetime_default(monkeypatch):
+    """Test that _state_cookie_lifetime defaults to 180 seconds (3 minutes)."""
+    monkeypatch.delenv("CHAINLIT_STATE_COOKIE_LIFETIME", raising=False)
+    importlib.reload(cookie_module)
+    assert cookie_module._state_cookie_lifetime == 180
+
+
+def test_state_cookie_lifetime_custom(monkeypatch):
+    """Test that _state_cookie_lifetime can be set via environment variable."""
+    monkeypatch.setenv("CHAINLIT_STATE_COOKIE_LIFETIME", "600")
+    importlib.reload(cookie_module)
+    assert cookie_module._state_cookie_lifetime == 600
+
+
+def test_clear_auth_cookie(client):
+    """Test cookie clearing removes all chunks."""
+    # Set initial token
+    client.post("/set-cookie", data={"token": "x" * 4000})
+
+    # Verify cookies exist
+    assert len(client.cookies) > 0
+
+    # Clear cookies
+    clear_response = client.delete("/clear-cookie")
+    assert clear_response.status_code == 200
+
+    # Verify cookies were cleared
+    assert len(clear_response.cookies) == 0
+    final_response = client.get("/get-token")
+    assert final_response.json()["token"] is None