refactor

apps/client/public/locales/en-US/translation.json

@@ -623,7 +623,7 @@
   "Upgrade your plan": "Upgrade your plan",
   "Available with a paid license": "Available with a paid license",
   "Upgrade your license tier.": "Upgrade your license tier.",
-  "AI features require a paid plan. Visit docmost.com for more information.": "AI features require a paid plan. Visit docmost.com for more information.",
+  "AI is only available in the Docmost enterprise edition. Contact sales@docmost.com.": "AI is only available in the Docmost enterprise edition. Contact sales@docmost.com.",
   "AI & MCP": "AI & MCP",
   "AI": "AI",
   "MCP": "MCP",

apps/client/src/components/settings/settings-sidebar.tsx

@@ -21,7 +21,7 @@ import { useTranslation } from "react-i18next";
 import { isCloud } from "@/lib/config.ts";
 import useUserRole from "@/hooks/use-user-role.tsx";
 import { useAtom } from "jotai";
-import { workspaceAtom } from "@/features/user/atoms/current-user-atom";
+import { entitlementAtom } from "@/ee/entitlement/entitlement-atom";
 import { Feature } from "@/ee/features";
 import { useUpgradeLabel } from "@/ee/hooks/use-upgrade-label";
 import {
@@ -136,7 +136,7 @@ export default function SettingsSidebar() {
   const [active, setActive] = useState(location.pathname);
   const { goBack } = useSettingsNavigation();
   const { isAdmin, isOwner } = useUserRole();
-  const [workspace] = useAtom(workspaceAtom);
+  const [entitlements] = useAtom(entitlementAtom);
   const upgradeLabel = useUpgradeLabel();
   const [mobileSidebarOpened] = useAtom(mobileSidebarAtom);
   const toggleMobileSidebar = useToggleSidebar(mobileSidebarAtom);
@@ -146,7 +146,7 @@ export default function SettingsSidebar() {
   }, [location.pathname]);
 
   const hasFeature = (f: string) =>
-    workspace?.features?.includes(f) ?? false;
+    entitlements?.features?.includes(f) ?? false;
 
   const canShowItem = (item: DataItem) => {
     if (item.env === "cloud" && !isCloud()) return false;
@@ -191,7 +191,7 @@ export default function SettingsSidebar() {
               prefetchHandler = prefetchBilling;
               break;
             case "License & Edition":
-              if (workspace?.hasLicenseKey) {
+              if (entitlements?.tier !== "free") {
                 prefetchHandler = prefetchLicense;
               }
               break;

apps/client/src/ee/ai/pages/ai-settings.tsx

@@ -63,7 +63,7 @@ export default function AiSettings() {
               mb="lg"
             >
               {t(
-                "AI features require a paid plan. Visit docmost.com for more information.",
+                "AI is only available in the Docmost enterprise edition. Contact sales@docmost.com.",
               )}
             </Alert>
           )}

apps/client/src/ee/components/sso-login.tsx

@@ -56,7 +56,7 @@ export default function SsoLogin() {
         />
       )}
 
-      {(data.features?.length > 0) && (
+      {data.authProviders.length > 0 && (
         <>
           <Stack align="stretch" justify="center" gap="sm">
             {data.authProviders.map((provider) => (

apps/client/src/ee/entitlement/entitlement-atom.ts

@@ -0,0 +1,5 @@
+import { atom } from "jotai";
+import type { Entitlements } from "./entitlement.types";
+
+const initialValue: Entitlements | null = null;
+export const entitlementAtom = atom(initialValue);

apps/client/src/ee/entitlement/entitlement-service.ts

@@ -0,0 +1,7 @@
+import api from "@/lib/api-client";
+import { Entitlements } from "./entitlement.types";
+
+export async function getEntitlements(): Promise<Entitlements> {
+  const req = await api.post<Entitlements>("/workspace/entitlements");
+  return req.data as Entitlements;
+}

apps/client/src/ee/entitlement/entitlement.types.ts

@@ -0,0 +1,7 @@
+export type Tier = "free" | "standard" | "business" | "enterprise";
+
+export type Entitlements = {
+  cloud: boolean;
+  tier: Tier;
+  features: string[];
+};

apps/client/src/ee/entitlement/use-entitlements.ts

@@ -0,0 +1,11 @@
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { getEntitlements } from "./entitlement-service";
+import { Entitlements } from "./entitlement.types";
+
+export function useEntitlements(): UseQueryResult<Entitlements> {
+  return useQuery({
+    queryKey: ["entitlements"],
+    queryFn: getEntitlements,
+    staleTime: 5 * 60 * 1000,
+  });
+}

apps/client/src/ee/hooks/use-feature.ts

@@ -1,12 +1,7 @@
 import { useAtom } from "jotai";
-import { workspaceAtom } from "@/features/user/atoms/current-user-atom";
+import { entitlementAtom } from "@/ee/entitlement/entitlement-atom";
 
 export const useHasFeature = (feature: string): boolean => {
-  const [workspace] = useAtom(workspaceAtom);
+  const [entitlements] = useAtom(entitlementAtom);
-  return workspace?.features?.includes(feature) ?? false;
+  return entitlements?.features?.includes(feature) ?? false;
-};
-
-export const useHasAnyFeature = (): boolean => {
-  const [workspace] = useAtom(workspaceAtom);
-  return (workspace?.features?.length ?? 0) > 0;
 };

apps/client/src/ee/hooks/use-upgrade-label.ts

@@ -1,14 +1,14 @@
 import { useAtom } from "jotai";
 import { useTranslation } from "react-i18next";
-import { workspaceAtom } from "@/features/user/atoms/current-user-atom";
+import { entitlementAtom } from "@/ee/entitlement/entitlement-atom";
 import { isCloud } from "@/lib/config";
 
 export function useUpgradeLabel(): string {
   const { t } = useTranslation();
-  const [workspace] = useAtom(workspaceAtom);
+  const [entitlements] = useAtom(entitlementAtom);
 
   if (!isCloud()) {
-    return workspace?.hasLicenseKey
+    return entitlements != null && entitlements.tier !== "free"
       ? t("Upgrade your license tier.")
       : t("Available with a paid license");
   }

apps/client/src/ee/licence/components/activate-license-modal.tsx

@@ -7,21 +7,22 @@ import { useTranslation } from "react-i18next";
 import { useActivateMutation } from "@/ee/licence/queries/license-query.ts";
 import { useDisclosure } from "@mantine/hooks";
 import { useAtom } from "jotai";
-import { workspaceAtom } from "@/features/user/atoms/current-user-atom.ts";
+import { entitlementAtom } from "@/ee/entitlement/entitlement-atom";
 import RemoveLicense from "@/ee/licence/components/remove-license.tsx";
 
 export default function ActivateLicense() {
   const { t } = useTranslation();
   const [opened, { open, close }] = useDisclosure(false);
-  const [workspace] = useAtom(workspaceAtom);
+  const [entitlements] = useAtom(entitlementAtom);
+  const hasLicense = entitlements != null && entitlements.tier !== "free";
 
   return (
     <Group justify="flex-end" wrap="nowrap" mb="sm">
       <Button onClick={open}>
-        {workspace?.hasLicenseKey ? t("Update license") : t("Add license")}
+        {hasLicense ? t("Update license") : t("Add license")}
       </Button>
 
-      {workspace?.hasLicenseKey && <RemoveLicense />}
+      {hasLicense && <RemoveLicense />}
 
       <Modal
         size="550"
@@ -59,7 +60,7 @@ export function ActivateLicenseForm({ onClose }: ActivateLicenseFormProps) {
   async function handleSubmit(data: { licenseKey: string }) {
     await activateLicenseMutation.mutateAsync(data.licenseKey);
     form.reset();
-    onClose();
+    onClose?.();
   }
 
   return (

apps/client/src/ee/licence/pages/license.tsx

@@ -8,10 +8,11 @@ import ActivateLicenseForm from "@/ee/licence/components/activate-license-modal.
 import InstallationDetails from "@/ee/licence/components/installation-details.tsx";
 import OssDetails from "@/ee/licence/components/oss-details.tsx";
 import { useAtom } from "jotai/index";
-import { workspaceAtom } from "@/features/user/atoms/current-user-atom.ts";
+import { entitlementAtom } from "@/ee/entitlement/entitlement-atom";
 
 export default function License() {
-  const [workspace] = useAtom(workspaceAtom);
+  const [entitlements] = useAtom(entitlementAtom);
+  const hasLicense = entitlements != null && entitlements.tier !== "free";
   const { isAdmin } = useUserRole();
 
   if (!isAdmin) {
@@ -29,7 +30,7 @@ export default function License() {
 
       <InstallationDetails />
 
-      {workspace?.hasLicenseKey ? <LicenseDetails /> : <OssDetails />}
+      {hasLicense ? <LicenseDetails /> : <OssDetails />}
     </>
   );
 }

apps/client/src/ee/licence/queries/license-query.ts

@@ -31,6 +31,7 @@ export function useActivateMutation() {
         queryKey: ["license"],
       });
       queryClient.refetchQueries({ queryKey: ["currentUser"] });
+      queryClient.refetchQueries({ queryKey: ["entitlements"] });
     },
     onError: (error) => {
       const errorMessage = error["response"]?.data?.message;
@@ -47,6 +48,7 @@ export function useRemoveLicenseMutation() {
     onSuccess: () => {
       queryClient.refetchQueries({ queryKey: ["license"] });
       queryClient.refetchQueries({ queryKey: ["currentUser"] });
+      queryClient.refetchQueries({ queryKey: ["entitlements"] });
     },
   });
 }

apps/client/src/features/user/user-provider.tsx

@@ -1,4 +1,4 @@
-import { useAtom } from "jotai";
+import { useAtom, useSetAtom } from "jotai";
 import { currentUserAtom } from "@/features/user/atoms/current-user-atom";
 import React, { useEffect } from "react";
 import useCurrentUser from "@/features/user/hooks/use-current-user";
@@ -11,10 +11,14 @@ import { useTreeSocket } from "@/features/websocket/use-tree-socket.ts";
 import { useNotificationSocket } from "@/features/notification/hooks/use-notification-socket.ts";
 import { useCollabToken } from "@/features/auth/queries/auth-query.tsx";
 import { Error404 } from "@/components/ui/error-404.tsx";
+import { useEntitlements } from "@/ee/entitlement/use-entitlements";
+import { entitlementAtom } from "@/ee/entitlement/entitlement-atom";
 
 export function UserProvider({ children }: React.PropsWithChildren) {
   const [, setCurrentUser] = useAtom(currentUserAtom);
+  const setEntitlements = useSetAtom(entitlementAtom);
   const { data, isLoading, error, isError } = useCurrentUser();
+  const { data: entitlements } = useEntitlements();
   const { i18n } = useTranslation();
   const [, setSocket] = useAtom(socketAtom);
   // fetch collab token on load
@@ -56,6 +60,12 @@ export function UserProvider({ children }: React.PropsWithChildren) {
     }
   }, [data, isLoading]);
 
+  useEffect(() => {
+    if (entitlements) {
+      setEntitlements(entitlements);
+    }
+  }, [entitlements]);
+
   if (isLoading) return <></>;
 
   if (isError && error?.["response"]?.status === 404) {

apps/client/src/features/workspace/types/workspace.types.ts

@@ -20,8 +20,6 @@ export interface IWorkspace {
   emailDomains: string[];
   memberCount?: number;
   plan?: string;
-  hasLicenseKey?: boolean;
-  features?: string[];
   enforceMfa?: boolean;
   aiSearch?: boolean;
   generativeAi?: boolean;
@@ -85,7 +83,6 @@ export interface IPublicWorkspace {
   hostname: string;
   enforceSso: boolean;
   authProviders: IAuthProvider[];
-  features?: string[];
 }
 
 export interface IVersion {

apps/server/src/core/user/user.controller.ts

@@ -13,7 +13,6 @@ import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
 import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
 import { User, Workspace } from '@docmost/db/types/entity.types';
 import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
-import { LicenseCheckService } from '../../integrations/environment/license-check.service';
 
 @UseGuards(JwtAuthGuard)
 @Controller('users')
@@ -21,7 +20,6 @@ export class UserController {
   constructor(
     private readonly userService: UserService,
     private readonly workspaceRepo: WorkspaceRepo,
-    private readonly licenseCheckService: LicenseCheckService,
   ) {}
 
   @HttpCode(HttpStatus.OK)
@@ -36,16 +34,9 @@ export class UserController {
 
     const { licenseKey, ...rest } = workspace;
 
-    const features = this.licenseCheckService.resolveFeatures(
-      licenseKey,
-      rest.plan,
-    );
-
     const workspaceInfo = {
       ...rest,
       memberCount,
-      hasLicenseKey: Boolean(licenseKey),
-      features,
     };
 
     return { user: authUser, workspace: workspaceInfo };

apps/server/src/core/workspace/controllers/workspace.controller.ts

@@ -32,8 +32,10 @@ import {
 } from '../../casl/interfaces/workspace-ability.type';
 import { FastifyReply } from 'fastify';
 import { EnvironmentService } from '../../../integrations/environment/environment.service';
+import { LicenseCheckService } from '../../../integrations/environment/license-check.service';
 import { CheckHostnameDto } from '../dto/check-hostname.dto';
 import { RemoveWorkspaceUserDto } from '../dto/remove-workspace-user.dto';
+import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
 
 @UseGuards(JwtAuthGuard)
 @Controller('workspace')
@@ -42,7 +44,9 @@ export class WorkspaceController {
     private readonly workspaceService: WorkspaceService,
     private readonly workspaceInvitationService: WorkspaceInvitationService,
     private readonly workspaceAbility: WorkspaceAbilityFactory,
+    private readonly workspaceRepo: WorkspaceRepo,
     private environmentService: EnvironmentService,
+    private licenseCheckService: LicenseCheckService,
   ) {}
 
   @Public()
@@ -58,6 +62,23 @@ export class WorkspaceController {
     return this.workspaceService.getWorkspaceInfo(workspace.id);
   }
 
+  @HttpCode(HttpStatus.OK)
+  @Post('entitlements')
+  async getEntitlements(@AuthWorkspace() workspace: Workspace) {
+    let { licenseKey } = workspace;
+    const { plan } = workspace;
+
+    if (!licenseKey) {
+      licenseKey = await this.workspaceRepo.findLicenseKeyById(workspace.id);
+    }
+
+    return {
+      cloud: this.environmentService.isCloud(),
+      tier: this.licenseCheckService.resolveTier(licenseKey, plan),
+      features: this.licenseCheckService.resolveFeatures(licenseKey, plan),
+    };
+  }
+
   @HttpCode(HttpStatus.OK)
   @Post('update')
   async updateWorkspace(

apps/server/src/core/workspace/services/workspace.service.ts

@@ -106,13 +106,9 @@ export class WorkspaceService {
       throw new NotFoundException('Workspace not found');
     }
 
-    const { licenseKey, ...rest } = workspace;
+    const { licenseKey, plan, ...rest } = workspace;
 
-    return {
+    return rest;
-      ...rest,
-      hasLicenseKey: Boolean(licenseKey),
-      features: this.licenseCheckService.resolveFeatures(licenseKey, rest.plan),
-    };
   }
 
   async create(
@@ -337,6 +333,10 @@ export class WorkspaceService {
         .where('id', '=', workspaceId)
         .executeTakeFirst();
 
+      if (!ws) {
+        throw new NotFoundException('Workspace not found');
+      }
+
       if (!this.licenseCheckService.hasFeature(ws.licenseKey, 'security:settings')) {
         throw new ForbiddenException(
           'This feature requires a valid license',
@@ -504,11 +504,7 @@ export class WorkspaceService {
     }
 
     const { licenseKey, ...rest } = workspace;
-    return {
+    return rest;
-      ...rest,
-      hasLicenseKey: Boolean(licenseKey),
-      features: this.licenseCheckService.resolveFeatures(licenseKey, rest.plan),
-    };
   }
 
   async getWorkspaceUsers(

apps/server/src/integrations/environment/license-check.service.ts

@@ -69,4 +69,25 @@ export class LicenseCheckService {
 
     return this.getFeatures(licenseKey);
   }
+
+  resolveTier(licenseKey: string, plan: string): string {
+    if (this.environmentService.isCloud()) {
+      return plan ?? 'standard';
+    }
+
+    return this.getLicenseType(licenseKey) ?? 'free';
+  }
+
+  private getLicenseType(licenseKey: string): string | null {
+    try {
+      // eslint-disable-next-line @typescript-eslint/no-require-imports
+      const LicenseModule = require('../../ee/licence/license.service');
+      const licenseService = this.moduleRef.get(LicenseModule.LicenseService, {
+        strict: false,
+      });
+      return licenseService.getLicenseType(licenseKey);
+    } catch {
+      return null;
+    }
+  }
 }