From d48447a352ddec3d4ad4d8ee5a169529888f23b2 Mon Sep 17 00:00:00 2001
From: denis0001-dev <denis0001.dev@ya.ru>
Date: Mon, 29 Sep 2025 17:39:20 +0300
Subject: [PATCH] Make registration immediately log the user in and create the
 encryption keys

---
 backend/routes/account.py                 |  6 +++++-
 frontend/src/ui/screen/RegisterScreen.tsx | 21 +++++++++++++++------
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/backend/routes/account.py b/backend/routes/account.py
index 2d12a0f..4f0e0af 100644
--- a/backend/routes/account.py
+++ b/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")
diff --git a/frontend/src/ui/screen/RegisterScreen.tsx b/frontend/src/ui/screen/RegisterScreen.tsx
index b03c35f..a93efd3 100644
--- a/frontend/src/ui/screen/RegisterScreen.tsx
+++ b/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
-                            showAlert("success", "Регистрация прошла успешно! Теперь вы можете войти.");
-                            await delay(2000);
-                            setCurrentPage("login");
+                            const data: LoginResponse = await response.json();
+                            // Store the JWT token first
+                            setUser(data.token, data.user);
+                            
+                            // 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 || "Ошибка при регистрации");