From a70018f70afc42ad86bb357c51ed13c48cfea03e Mon Sep 17 00:00:00 2001
From: Mauricio Siu <siumauricio@icloud.com>
Date: Sun, 1 Feb 2026 22:01:13 -0600
Subject: [PATCH] feat(auth): add enterprise feature flags to user context and
 request validation

- Updated user context to include `enableEnterpriseFeatures` and `isValidEnterpriseLicense` properties.
- Modified request validation to set these properties based on user data, enhancing enterprise feature management.
- Adjusted the enterprise procedure to check user flags directly from the context instead of querying the database.
---
 apps/dokploy/server/api/trpc.ts | 23 ++++++++++-------------
 packages/server/src/lib/auth.ts | 17 +++++++++++++++++
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/apps/dokploy/server/api/trpc.ts b/apps/dokploy/server/api/trpc.ts
index ce8d8c4ea..51f8cdbee 100644
--- a/apps/dokploy/server/api/trpc.ts
+++ b/apps/dokploy/server/api/trpc.ts
@@ -7,10 +7,8 @@
  * need to use are documented accordingly near the end.
  */
 
-import { user as userSchema } from "@dokploy/server/db/schema";
 import { validateRequest } from "@dokploy/server/lib/auth";
 import type { OpenApiMeta } from "@dokploy/trpc-openapi";
-import { eq } from "drizzle-orm";
 import { initTRPC, TRPCError } from "@trpc/server";
 import type { CreateNextContextOptions } from "@trpc/server/adapters/next";
 import {
@@ -33,7 +31,14 @@ import { db } from "@/server/db";
  */
 
 interface CreateContextOptions {
-	user: (User & { role: "member" | "admin" | "owner"; ownerId: string }) | null;
+	user:
+		| (User & {
+				role: "member" | "admin" | "owner";
+				ownerId: string;
+				enableEnterpriseFeatures: boolean;
+				isValidEnterpriseLicense: boolean;
+		  })
+		| null;
 	session:
 		| (Session & { activeOrganizationId: string; impersonatedBy?: string })
 		| null;
@@ -234,17 +239,9 @@ export const enterpriseProcedure = t.procedure.use(async ({ ctx, next }) => {
 		throw new TRPCError({ code: "UNAUTHORIZED" });
 	}
 
-	const currentUser = await ctx.db.query.user.findFirst({
-		where: eq(userSchema.id, ctx.user.id),
-		columns: {
-			enableEnterpriseFeatures: true,
-			isValidEnterpriseLicense: true,
-		},
-	});
-
 	if (
-		!currentUser?.enableEnterpriseFeatures ||
-		!currentUser.isValidEnterpriseLicense
+		!ctx.user?.enableEnterpriseFeatures ||
+		!ctx.user.isValidEnterpriseLicense
 	) {
 		throw new TRPCError({
 			code: "FORBIDDEN",
diff --git a/packages/server/src/lib/auth.ts b/packages/server/src/lib/auth.ts
index 924467319..467ba54cb 100644
--- a/packages/server/src/lib/auth.ts
+++ b/packages/server/src/lib/auth.ts
@@ -279,6 +279,16 @@ const { handler, api } = betterAuth({
 				input: true,
 				defaultValue: "",
 			},
+			enableEnterpriseFeatures: {
+				type: "boolean",
+				required: false,
+				input: false,
+			},
+			isValidEnterpriseLicense: {
+				type: "boolean",
+				required: false,
+				input: false,
+			},
 		},
 	},
 	plugins: [
@@ -399,6 +409,8 @@ export const validateRequest = async (request: IncomingMessage) => {
 					twoFactorEnabled: userFromDb.twoFactorEnabled,
 					role: member?.role || "member",
 					ownerId: member?.organization.ownerId || apiKeyRecord.user.id,
+					enableEnterpriseFeatures: userFromDb.enableEnterpriseFeatures,
+					isValidEnterpriseLicense: userFromDb.isValidEnterpriseLicense,
 				},
 			};
 
@@ -437,10 +449,15 @@ export const validateRequest = async (request: IncomingMessage) => {
 			),
 			with: {
 				organization: true,
+				user: true,
 			},
 		});
 
 		session.user.role = member?.role || "member";
+		session.user.enableEnterpriseFeatures =
+			member?.user.enableEnterpriseFeatures || false;
+		session.user.isValidEnterpriseLicense =
+			member?.user.isValidEnterpriseLicense || false;
 		if (member) {
 			session.user.ownerId = member.organization.ownerId;
 		} else {