From 5e297f327a5e0c4a784cd2d4fb4e67636f9c2fda Mon Sep 17 00:00:00 2001
From: Hakan Shehu <hakanshehu15@gmail.com>
Date: Sun, 27 Oct 2024 16:19:24 +0100
Subject: [PATCH] Fix login/register redirect flow

---
 .../src/main/handlers/mutations/email-login.ts  |  2 ++
 .../main/handlers/mutations/email-register.ts   |  2 ++
 desktop/src/operations/mutations/email-login.ts | 17 +++++++++++++++--
 .../src/operations/mutations/email-register.ts  | 17 +++++++++++++++--
 .../components/accounts/email-login.tsx         | 16 ++++++++++++++--
 .../components/accounts/email-register.tsx      | 16 ++++++++++++++--
 .../components/workspaces/workspace-create.tsx  |  2 +-
 7 files changed, 63 insertions(+), 9 deletions(-)

diff --git a/desktop/src/main/handlers/mutations/email-login.ts b/desktop/src/main/handlers/mutations/email-login.ts
index 2c75dd91..d37ddcf1 100644
--- a/desktop/src/main/handlers/mutations/email-login.ts
+++ b/desktop/src/main/handlers/mutations/email-login.ts
@@ -120,6 +120,8 @@ export class EmailLoginMutationHandler
     return {
       output: {
         success: true,
+        account: data.account,
+        workspaces: data.workspaces,
       },
       changes: changedTables,
     };
diff --git a/desktop/src/main/handlers/mutations/email-register.ts b/desktop/src/main/handlers/mutations/email-register.ts
index c6b77fde..04df61b4 100644
--- a/desktop/src/main/handlers/mutations/email-register.ts
+++ b/desktop/src/main/handlers/mutations/email-register.ts
@@ -121,6 +121,8 @@ export class EmailRegisterMutationHandler
     return {
       output: {
         success: true,
+        account: data.account,
+        workspaces: data.workspaces,
       },
       changes: changedTables,
     };
diff --git a/desktop/src/operations/mutations/email-login.ts b/desktop/src/operations/mutations/email-login.ts
index c4ffb853..3a5d63a0 100644
--- a/desktop/src/operations/mutations/email-login.ts
+++ b/desktop/src/operations/mutations/email-login.ts
@@ -1,3 +1,6 @@
+import { Account } from '@/types/accounts';
+import { WorkspaceOutput } from '@/types/workspaces';
+
 export type EmailLoginMutationInput = {
   type: 'email_login';
   server: string;
@@ -5,8 +8,18 @@ export type EmailLoginMutationInput = {
   password: string;
 };
 
-export type EmailLoginMutationOutput = {
-  success: boolean;
+export type EmailLoginMutationOutput =
+  | EmailLoginMutationSuccessOutput
+  | EmailLoginMutationErrorOutput;
+
+export type EmailLoginMutationSuccessOutput = {
+  success: true;
+  account: Account;
+  workspaces: WorkspaceOutput[];
+};
+
+export type EmailLoginMutationErrorOutput = {
+  success: false;
 };
 
 declare module '@/operations/mutations' {
diff --git a/desktop/src/operations/mutations/email-register.ts b/desktop/src/operations/mutations/email-register.ts
index 9851671c..c0e485b2 100644
--- a/desktop/src/operations/mutations/email-register.ts
+++ b/desktop/src/operations/mutations/email-register.ts
@@ -1,3 +1,6 @@
+import { Account } from '@/types/accounts';
+import { WorkspaceOutput } from '@/types/workspaces';
+
 export type EmailRegisterMutationInput = {
   type: 'email_register';
   server: string;
@@ -6,8 +9,18 @@ export type EmailRegisterMutationInput = {
   password: string;
 };
 
-export type EmailRegisterMutationOutput = {
-  success: boolean;
+export type EmailRegisterMutationOutput =
+  | EmailRegisterMutationSuccessOutput
+  | EmailRegisterMutationErrorOutput;
+
+export type EmailRegisterMutationSuccessOutput = {
+  success: true;
+  account: Account;
+  workspaces: WorkspaceOutput[];
+};
+
+export type EmailRegisterMutationErrorOutput = {
+  success: false;
 };
 
 declare module '@/operations/mutations' {
diff --git a/desktop/src/renderer/components/accounts/email-login.tsx b/desktop/src/renderer/components/accounts/email-login.tsx
index e8133d40..d25b02d7 100644
--- a/desktop/src/renderer/components/accounts/email-login.tsx
+++ b/desktop/src/renderer/components/accounts/email-login.tsx
@@ -46,8 +46,20 @@ export const EmailLogin = ({ server }: EmailLoginProps) => {
         password: values.password,
         server: server.domain,
       },
-      onSuccess() {
-        navigate('/');
+      onSuccess(output) {
+        if (output.success) {
+          if (output.workspaces.length > 0) {
+            navigate(`/${output.workspaces[0].id}`);
+          } else {
+            navigate('/create');
+          }
+        } else {
+          toast({
+            title: 'Failed to login',
+            description: 'Invalid email or password.',
+            variant: 'destructive',
+          });
+        }
       },
       onError() {
         toast({
diff --git a/desktop/src/renderer/components/accounts/email-register.tsx b/desktop/src/renderer/components/accounts/email-register.tsx
index 3ec0307e..9656e207 100644
--- a/desktop/src/renderer/components/accounts/email-register.tsx
+++ b/desktop/src/renderer/components/accounts/email-register.tsx
@@ -49,8 +49,20 @@ export const EmailRegister = ({ server }: EmailRegisterProps) => {
         password: values.password,
         server: server.domain,
       },
-      onSuccess() {
-        navigate('/');
+      onSuccess(output) {
+        if (output.success) {
+          if (output.workspaces.length > 0) {
+            navigate(`/${output.workspaces[0].id}`);
+          } else {
+            navigate('/create');
+          }
+        } else {
+          toast({
+            title: 'Failed to register',
+            description: 'Email already in use.',
+            variant: 'destructive',
+          });
+        }
       },
       onError() {
         toast({
diff --git a/desktop/src/renderer/components/workspaces/workspace-create.tsx b/desktop/src/renderer/components/workspaces/workspace-create.tsx
index c4d303e3..7296f008 100644
--- a/desktop/src/renderer/components/workspaces/workspace-create.tsx
+++ b/desktop/src/renderer/components/workspaces/workspace-create.tsx
@@ -12,7 +12,7 @@ export const WorkspaceCreate = () => {
   const { mutate, isPending } = useMutation();
 
   const handleCancel =
-    account.workspaces.length > 1 ? () => navigate('/') : undefined;
+    account.workspaces.length > 0 ? () => navigate('/') : undefined;
 
   return (
     <div className="container flex flex-row justify-center">