handle join mechanism with cookie

Author: cyclotruc

src/backend/main.py

@@ -76,10 +76,21 @@ async def lifespan(app: FastAPI):
 app.mount("/assets", StaticFiles(directory=ASSETS_DIR), name="assets")
 app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 
-async def serve_index_html(request: Request = None):
+async def serve_index_html(request: Request = None, response: Response = None, pad_id: Optional[UUID] = None):
     """
     Helper function to serve the index.html file or proxy to dev server based on PAD_DEV_MODE.
+    Optionally sets a pending_pad_id cookie if pad_id is provided.
     """
+    # Set cookie if pad_id is provided
+    if pad_id is not None and response is not None:
+        response.set_cookie(
+            key="pending_pad_id",
+            value=str(pad_id),
+            httponly=True,
+            secure=True,
+            samesite="lax"
+        )
+    
     if PAD_DEV_MODE:
         try:
             # Proxy the request to the development server's root URL
@@ -89,72 +100,46 @@ async def serve_index_html(request: Request = None):
                 url = f"{DEV_FRONTEND_URL}{request.url.path}"
 
             async with httpx.AsyncClient() as client:
-                response = await client.get(url)
-                return Response(
-                    content=response.content, 
-                    status_code=response.status_code,
-                    media_type=response.headers.get("content-type")
-                )
+                proxy_response = await client.get(url)
+                response.body = proxy_response.content
+                response.status_code = proxy_response.status_code
+                response.media_type = proxy_response.headers.get("content-type")
+                return response
         except Exception as e:
             error_message = f"Error proxying to dev server: {e}"
             print(error_message)
-            return Response(
-                status_code=500,
-            )
+            response.status_code = 500
+            return response
     else:
         # For production, serve the static build
-        return FileResponse(os.path.join(STATIC_DIR, "index.html"))
+        return FileResponse(os.path.join(STATIC_DIR, "index.html"), background=response)
 
 @app.get("/pad/{pad_id}")
 async def read_pad(
     pad_id: UUID,
     request: Request,
+    response: Response,
     user: Optional[UserSession] = Depends(optional_auth),
     session: AsyncSession = Depends(get_session)
 ):
     if not user:
-        print("No user found")
-        return await serve_index_html(request)
+        return await serve_index_html(request, response, pad_id)
 
     try:
         pad = await Pad.get_by_id(session, pad_id)
         if not pad:
             print("No pad found")
-            return await serve_index_html(request)
+            return await serve_index_html(request, response)
 
-        # Check access permissions
         if not pad.can_access(user.id):
             print("No access to pad")
-            return await serve_index_html(request)
-            
-        # Worker assignment is now handled automatically in Pad.get_by_id()
-        print(f"Pad {pad_id} accessed by user {user.id}, worker: {pad.worker_id[:8] if pad.worker_id else 'None'}")
+            return await serve_index_html(request, response)
 
-        # Add pad to user's open pads if not already there and update last selected pad
-        user_obj = await User.get_by_id(session, user.id)
-        if user_obj:
-            # Convert all UUIDs to strings for comparison
-            open_pads_str = [str(pid) for pid in user_obj._store.open_pads]
-            if str(pad_id) not in open_pads_str:
-                # Convert back to UUIDs for storage
-                user_obj._store.open_pads = [UUID(pid) for pid in open_pads_str] + [pad_id]
-                try:
-                    await user_obj.save(session)
-                except Exception as e:
-                    print(f"Error updating user's open pads: {e}")
-                    # Continue even if update fails - don't block pad access
-            
-            # Update last selected pad
-            try:
-                await user_obj.set_last_selected_pad(session, pad_id)
-            except Exception as e:
-                print(f"Error updating last selected pad: {e}")
-                # Continue even if update fails - don't block pad access
-        
-        return await serve_index_html(request)
+        # Just serve the page if user has access
+        return await serve_index_html(request, response, pad_id)
     except Exception as e:
         print(f"Error in read_pad endpoint: {e}")
-        return await serve_index_html(request)
+        return await serve_index_html(request, response, pad_id)
 
 @app.get("/")
 async def read_root(request: Request, auth: Optional[UserSession] = Depends(optional_auth)):

src/backend/routers/users_router.py

@@ -4,20 +4,23 @@
 
 import posthog
 import jwt
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, Depends, HTTPException, Request, Response
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from cache import RedisClient
 from config import get_jwks_client, OIDC_CLIENT_ID, FRONTEND_URL
 from dependencies import UserSession, require_admin, require_auth
 from database.database import get_session
 from domain.user import User
+from domain.pad import Pad
 
 users_router = APIRouter()
 
 
 @users_router.get("/me")
 async def get_user_info(
+    request: Request,
+    response: Response,
     user: UserSession = Depends(require_auth),
     session: AsyncSession = Depends(get_session),
 ):
@@ -26,6 +29,34 @@ async def get_user_info(
     # Get the user from database to access last_selected_pad
     user_obj = await User.get_by_id(session, user.id)
 
+    # Check for pending pad cookie
+    pending_pad_id = request.cookies.get("pending_pad_id")
+    if pending_pad_id:
+        try:
+            pad_id = UUID(pending_pad_id)
+            pad = await Pad.get_by_id(session, pad_id)
+            
+            if pad and pad.can_access(user.id):
+                # Convert all UUIDs to strings for comparison
+                open_pads_str = [str(pid) for pid in user_obj._store.open_pads]
+                if str(pad_id) not in open_pads_str:
+                    user_obj._store.open_pads = [UUID(pid) for pid in open_pads_str] + [pad_id]
+                    await user_obj.save(session)
+                
+                # Update last selected pad
+                await user_obj.set_last_selected_pad(session, pad_id)
+            
+        except (ValueError, Exception) as e:
+            print(f"Error processing pending pad: {e}")
+        finally:
+            # Always clear the cookie with same settings as when setting it
+            response.delete_cookie(
+                key="pending_pad_id",
+                secure=True,
+                httponly=False,
+                samesite="lax"
+            )
+    
     # Create token data dictionary from UserSession properties
     token_data = {
         "id": user.id,