From 9809e054535436c759a9ac47267e12b90b7b7d72 Mon Sep 17 00:00:00 2001
From: Vasily Zubarev <me@vas3k.ru>
Date: Tue, 22 Apr 2025 14:43:45 +0200
Subject: [PATCH] fix: fetch session profile from DB

---
 components/settings/profile-settings-form.tsx |  6 ++--
 components/sidebar/sidebar-user.tsx           |  2 +-
 lib/auth.ts                                   | 30 +++++++++++++------
 models/users.ts                               |  6 ++++
 4 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/components/settings/profile-settings-form.tsx b/components/settings/profile-settings-form.tsx
index 9f91fdf..e45c9c9 100644
--- a/components/settings/profile-settings-form.tsx
+++ b/components/settings/profile-settings-form.tsx
@@ -34,10 +34,10 @@ export default function ProfileSettingsForm({ user }: { user: User }) {
       </form>
       <Card className="mt-4 p-4">
         <p>
-          Storage: {user.storageUsed ? formatBytes(user.storageUsed) : "N/A"} /{" "}
-          {user.storageLimit && user.storageLimit > 0 ? formatBytes(user.storageLimit) : "Unlimited"}
+          Storage Used: {formatBytes(user.storageUsed)} /{" "}
+          {user.storageLimit > 0 ? formatBytes(user.storageLimit) : "Unlimited"}
         </p>
-        <p>Tokens Balance: {user.tokenBalance ? formatNumber(user.tokenBalance) : "N/A"}</p>
+        <p>Tokens Balance: {formatNumber(user.tokenBalance)}</p>
       </Card>
     </div>
   )
diff --git a/components/sidebar/sidebar-user.tsx b/components/sidebar/sidebar-user.tsx
index b4ef6b7..4c88957 100644
--- a/components/sidebar/sidebar-user.tsx
+++ b/components/sidebar/sidebar-user.tsx
@@ -64,7 +64,7 @@ export default function SidebarUser({ profile, isSelfHosted }: { profile: UserPr
           <DropdownMenuItem asChild>
             <Link href="/settings/profile" className="flex items-center gap-2">
               <HardDrive className="h-4 w-4" />
-              Storage: {profile.storageUsed ? formatBytes(profile.storageUsed) : "N/A"}
+              Storage: {formatBytes(profile.storageUsed)}
             </Link>
           </DropdownMenuItem>
         </DropdownMenuGroup>
diff --git a/lib/auth.ts b/lib/auth.ts
index eda418a..aae418d 100644
--- a/lib/auth.ts
+++ b/lib/auth.ts
@@ -1,6 +1,6 @@
 import config from "@/lib/config"
 import { createUserDefaults } from "@/models/defaults"
-import { getSelfHostedUser, getUserByEmail } from "@/models/users"
+import { getSelfHostedUser, getUserByEmail, getUserById } from "@/models/users"
 import { stripe } from "@better-auth/stripe"
 import { User } from "@prisma/client"
 import { betterAuth } from "better-auth"
@@ -19,8 +19,9 @@ export type UserProfile = {
   name: string
   email: string
   avatar?: string
-  storageUsed?: number
-  tokenBalance?: number
+  storageUsed: number
+  storageLimit: number
+  tokenBalance: number
 }
 
 export const auth = betterAuth({
@@ -91,13 +92,24 @@ export async function getSession() {
 }
 
 export async function getCurrentUser(): Promise<User> {
-  const session = await getSession()
-  if (!session || !session.user) {
-    if (config.selfHosted.isEnabled) {
-      redirect(config.selfHosted.redirectUrl)
+  if (config.selfHosted.isEnabled) {
+    const user = await getSelfHostedUser()
+    if (user) {
+      return user
     } else {
-      redirect(config.auth.loginUrl)
+      redirect(config.selfHosted.redirectUrl)
     }
   }
-  return session.user as User
+
+  // Try to return user from session
+  const session = await getSession()
+  if (session && session.user) {
+    const user = await getUserById(session.user.id)
+    if (user) {
+      return user
+    }
+  }
+
+  // No session or user found
+  redirect(config.auth.loginUrl)
 }
diff --git a/models/users.ts b/models/users.ts
index caee9be..caf61b2 100644
--- a/models/users.ts
+++ b/models/users.ts
@@ -21,6 +21,12 @@ export const createSelfHostedUser = cache(async () => {
   })
 })
 
+export const getUserById = cache(async (id: string) => {
+  return await prisma.user.findUnique({
+    where: { id },
+  })
+})
+
 export const getUserByEmail = cache(async (email: string) => {
   return await prisma.user.findUnique({
     where: { email },