Make registration immediately log the user in and create the encryption keys

backend/routes/account.py

@@ -116,9 +116,13 @@ def register(request: RegisterRequest, db: Session = Depends(get_db)):
     db.commit()
     db.refresh(new_user)
 
+    token = create_token(new_user.id, new_user.username)
+
     return {
         "status": "success",
-        "message": "Регистрация прошла успешно Теперь вы можете войти."
+        "message": "Регистрация прошла успешно",
+        "token": token,
+        "user": convert_user(new_user)
     }
 
 @router.get("/crypto/public-key")

frontend/src/ui/screen/RegisterScreen.tsx

@@ -4,15 +4,16 @@ import { AuthContainer, AuthHeader } from "../components/Auth";
 import { AlertsContainer, type Alert, type AlertType } from "../components/Alerts";
 import { useRef } from "react";
 import { TextField } from "mdui/components/text-field";
-import type { ErrorResponse, RegisterRequest } from "../../core/types";
+import type { ErrorResponse, RegisterRequest, LoginResponse } from "../../core/types";
 import { API_BASE_URL } from "../../core/config";
-import { delay } from "../../utils/utils";
 import { useAppState } from "../state";
 import { MaterialTextField } from "../components/core/TextField";
+import { ensureKeysOnLogin } from "../../auth/crypto";
 
 export default function RegisterScreen() {
     const [alerts, updateAlerts] = useImmer<Alert[]>([]);
     const setCurrentPage = useAppState(state => state.setCurrentPage);
+    const setUser = useAppState(state => state.setUser);
 
     function showAlert(type: AlertType, message: string) {
         updateAlerts((alerts) => { alerts.push({type: type, message: message}) });
@@ -71,10 +72,18 @@ export default function RegisterScreen() {
                         });
                         
                         if (response.ok) {
-                            // Registration successful
+                            const data: LoginResponse = await response.json();
-                            showAlert("success", "Регистрация прошла успешно! Теперь вы можете войти.");
+                            // Store the JWT token first
-                            await delay(2000);
+                            setUser(data.token, data.user);
-                            setCurrentPage("login");
+                            
+                            // Setup keys with the token we just received
+                            try {
+                                await ensureKeysOnLogin(password, data.token);
+                            } catch (e) {
+                                console.error("Key setup failed:", e);
+                            }
+
+                            setCurrentPage("chat");
                         } else {
                             const data: ErrorResponse = await response.json();
                             showAlert("danger", data.message || "Ошибка при регистрации");