Update main.py

main.py

@@ -343,7 +343,9 @@ async def register_post(
 
     try:
         user = UserCreate(username=username, email=email, password=password, confirm_password=confirm_password)
-        register(user, db)  # If this function raises an HTTPException, it should be handled
+        registered_user = register(user, db)
+        # If registration is successful, store user info in the session
+        request.session["user_info"] = {"username": registered_user.username, "email": registered_user.email}
     except HTTPException as e:
         # Return to the registration page with the error detail
         return templates.TemplateResponse("register.html", {