From c98e717f5b003e119408d05086da7db1acab853f Mon Sep 17 00:00:00 2001
From: johnnyjoy <yourselfhosted@gmail.com>
Date: Sat, 17 Aug 2024 19:22:25 +0800
Subject: [PATCH] chore: fix feature checks

---
 frontend/web/src/App.tsx                                 | 4 ++--
 frontend/web/src/components/Header.tsx                   | 6 +++---
 frontend/web/src/components/Logo.tsx                     | 4 ++--
 frontend/web/src/components/setting/WorkspaceSection.tsx | 3 +--
 frontend/web/src/pages/SubscriptionSetting.tsx           | 8 ++++----
 frontend/web/src/pages/WorkspaceSetting.tsx              | 3 +--
 frontend/web/src/stores/workspace.ts                     | 3 +++
 7 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/frontend/web/src/App.tsx b/frontend/web/src/App.tsx
index bc38fdd..721a747 100644
--- a/frontend/web/src/App.tsx
+++ b/frontend/web/src/App.tsx
@@ -4,7 +4,7 @@ import { Outlet } from "react-router-dom";
 import DemoBanner from "@/components/DemoBanner";
 import { useWorkspaceStore } from "@/stores";
 import useNavigateTo from "./hooks/useNavigateTo";
-import { PlanType } from "./types/proto/api/v1/subscription_service";
+import { FeatureType } from "./stores/workspace";
 
 function App() {
   const navigateTo = useNavigateTo();
@@ -28,7 +28,7 @@ function App() {
   }, [workspaceStore.setting.customStyle]);
 
   useEffect(() => {
-    const hasCustomBranding = workspaceStore.profile.subscription?.plan === PlanType.PRO;
+    const hasCustomBranding = workspaceStore.checkFeatureAvailable(FeatureType.CustomeBranding);
     if (!hasCustomBranding || !workspaceStore.setting.branding) {
       return;
     }
diff --git a/frontend/web/src/components/Header.tsx b/frontend/web/src/components/Header.tsx
index 11b26f3..34186f7 100644
--- a/frontend/web/src/components/Header.tsx
+++ b/frontend/web/src/components/Header.tsx
@@ -16,7 +16,7 @@ const Header: React.FC = () => {
   const workspaceStore = useWorkspaceStore();
   const currentUser = useUserStore().getCurrentUser();
   const [showAboutDialog, setShowAboutDialog] = useState<boolean>(false);
-  const profile = workspaceStore.profile;
+  const subscription = workspaceStore.getSubscription();
   const isAdmin = currentUser.role === Role.ADMIN;
   const shouldShowRouterSwitch = location.pathname === "/shortcuts" || location.pathname === "/collections";
   const selectedSection = location.pathname === "/shortcuts" ? "Shortcuts" : location.pathname === "/collections" ? "Collections" : "";
@@ -35,10 +35,10 @@ const Header: React.FC = () => {
               <Logo className="mr-2" />
               Slash
             </Link>
-            {profile.subscription?.plan && [PlanType.PRO, PlanType.ENTERPRISE].includes(profile.subscription.plan) && (
+            {[PlanType.PRO, PlanType.ENTERPRISE].includes(subscription.plan) && (
               <span className="ml-1 text-xs px-1.5 leading-5 border rounded-full bg-blue-600 border-blue-700 text-white shadow dark:opacity-70">
                 {/* PRO or ENT */}
-                {profile.subscription.plan.substring(0, 3)}
+                {subscription.plan.substring(0, 3)}
               </span>
             )}
             {shouldShowRouterSwitch && (
diff --git a/frontend/web/src/components/Logo.tsx b/frontend/web/src/components/Logo.tsx
index f8a69fd..c6e01b6 100644
--- a/frontend/web/src/components/Logo.tsx
+++ b/frontend/web/src/components/Logo.tsx
@@ -1,6 +1,6 @@
 import classNames from "classnames";
 import { useWorkspaceStore } from "@/stores";
-import { PlanType } from "@/types/proto/api/v1/subscription_service";
+import { FeatureType } from "@/stores/workspace";
 import Icon from "./Icon";
 
 interface Props {
@@ -9,7 +9,7 @@ interface Props {
 
 const Logo = ({ className }: Props) => {
   const workspaceStore = useWorkspaceStore();
-  const hasCustomBranding = workspaceStore.profile.subscription?.plan === PlanType.PRO;
+  const hasCustomBranding = workspaceStore.checkFeatureAvailable(FeatureType.CustomeBranding);
   const branding = hasCustomBranding && workspaceStore.setting.branding ? new TextDecoder().decode(workspaceStore.setting.branding) : "";
   return (
     <div className={classNames("w-8 h-auto dark:text-gray-500 rounded-lg overflow-hidden", className)}>
diff --git a/frontend/web/src/components/setting/WorkspaceSection.tsx b/frontend/web/src/components/setting/WorkspaceSection.tsx
index e09e53c..877181f 100644
--- a/frontend/web/src/components/setting/WorkspaceSection.tsx
+++ b/frontend/web/src/components/setting/WorkspaceSection.tsx
@@ -7,7 +7,6 @@ import { workspaceServiceClient } from "@/grpcweb";
 import { useWorkspaceStore } from "@/stores";
 import { FeatureType } from "@/stores/workspace";
 import { Visibility } from "@/types/proto/api/v1/common";
-import { PlanType } from "@/types/proto/api/v1/subscription_service";
 import { WorkspaceSetting } from "@/types/proto/api/v1/workspace_service";
 import FeatureBadge from "../FeatureBadge";
 import Icon from "../Icon";
@@ -35,7 +34,7 @@ const WorkspaceSection = () => {
   const [workspaceSetting, setWorkspaceSetting] = useState<WorkspaceSetting>(workspaceStore.setting);
   const originalWorkspaceSetting = useRef<WorkspaceSetting>(workspaceStore.setting);
   const allowSave = !isEqual(originalWorkspaceSetting.current, workspaceSetting);
-  const hasCustomBranding = workspaceStore.profile.subscription?.plan === PlanType.PRO;
+  const hasCustomBranding = workspaceStore.checkFeatureAvailable(FeatureType.CustomeBranding);
   const branding = hasCustomBranding && workspaceSetting.branding ? new TextDecoder().decode(workspaceSetting.branding) : "";
 
   const onBrandingChange = async (event: React.ChangeEvent<HTMLInputElement>) => {
diff --git a/frontend/web/src/pages/SubscriptionSetting.tsx b/frontend/web/src/pages/SubscriptionSetting.tsx
index 9ce38f0..6842cae 100644
--- a/frontend/web/src/pages/SubscriptionSetting.tsx
+++ b/frontend/web/src/pages/SubscriptionSetting.tsx
@@ -14,7 +14,7 @@ const SubscriptionSetting: React.FC = () => {
   const currentUser = useUserStore().getCurrentUser();
   const [licenseKey, setLicenseKey] = useState<string>("");
   const isAdmin = currentUser.role === Role.ADMIN;
-  const profile = workspaceStore.profile;
+  const subscription = workspaceStore.getSubscription();
 
   const handleDeleteLicenseKey = async () => {
     if (!isAdmin) {
@@ -58,7 +58,7 @@ const SubscriptionSetting: React.FC = () => {
         <p className="text-2xl shrink-0 font-semibold text-gray-900 dark:text-gray-500">Subscription</p>
         <div className="mt-2">
           <span className="text-gray-500 mr-2">Current plan:</span>
-          <span className="text-2xl mr-4 dark:text-gray-400">{stringifyPlanType(profile.subscription?.plan)}</span>
+          <span className="text-2xl mr-4 dark:text-gray-400">{stringifyPlanType(subscription.plan)}</span>
         </div>
         <Textarea
           className="w-full mt-2"
@@ -70,7 +70,7 @@ const SubscriptionSetting: React.FC = () => {
         />
         <div className="w-full flex justify-between items-center mt-4">
           <div>
-            {profile.subscription?.plan === PlanType.FREE && (
+            {subscription.plan === PlanType.FREE && (
               <Link href="https://yourselfhosted.lemonsqueezy.com/checkout/buy/947e9a56-c93a-4294-8d71-2ea4b0f3ec51" target="_blank">
                 Buy a license key
                 <Icon.ExternalLink className="w-4 h-auto ml-1" />
@@ -78,7 +78,7 @@ const SubscriptionSetting: React.FC = () => {
             )}
           </div>
           <div className="flex justify-end items-center gap-2">
-            {profile.subscription?.plan !== PlanType.FREE && (
+            {subscription.plan !== PlanType.FREE && (
               <Button color="neutral" variant="plain" onClick={handleDeleteLicenseKey}>
                 Reset
               </Button>
diff --git a/frontend/web/src/pages/WorkspaceSetting.tsx b/frontend/web/src/pages/WorkspaceSetting.tsx
index fc520bb..8a6cdae 100644
--- a/frontend/web/src/pages/WorkspaceSetting.tsx
+++ b/frontend/web/src/pages/WorkspaceSetting.tsx
@@ -12,7 +12,6 @@ const WorkspaceSetting = () => {
   const workspaceStore = useWorkspaceStore();
   const currentUser = useUserStore().getCurrentUser();
   const isAdmin = currentUser.role === Role.ADMIN;
-  const profile = workspaceStore.profile;
 
   useEffect(() => {
     if (!isAdmin) {
@@ -33,7 +32,7 @@ const WorkspaceSetting = () => {
         <p className="text-2xl shrink-0 font-semibold text-gray-900 dark:text-gray-500">Subscription</p>
         <div className="mt-2">
           <span className="text-gray-500 mr-2">Current plan:</span>
-          <span className="text-2xl mr-4 dark:text-gray-400">{stringifyPlanType(profile.subscription?.plan)}</span>
+          <span className="text-2xl mr-4 dark:text-gray-400">{stringifyPlanType(workspaceStore.getSubscription().plan)}</span>
           <Link to="/setting/subscription" unstable_viewTransition>
             <Button size="sm" variant="outlined" startDecorator={<Icon.Settings className="w-4 h-auto" />}>
               Manage
diff --git a/frontend/web/src/stores/workspace.ts b/frontend/web/src/stores/workspace.ts
index bc0f27d..b1f278e 100644
--- a/frontend/web/src/stores/workspace.ts
+++ b/frontend/web/src/stores/workspace.ts
@@ -1,5 +1,6 @@
 import { create } from "zustand";
 import { workspaceServiceClient } from "@/grpcweb";
+import { Subscription } from "@/types/proto/api/v1/subscription_service";
 import { WorkspaceProfile, WorkspaceSetting } from "@/types/proto/api/v1/workspace_service";
 
 export enum FeatureType {
@@ -17,6 +18,7 @@ interface WorkspaceState {
   // Workspace related actions.
   fetchWorkspaceProfile: () => Promise<WorkspaceProfile>;
   fetchWorkspaceSetting: () => Promise<WorkspaceSetting>;
+  getSubscription: () => Subscription;
   checkFeatureAvailable: (feature: FeatureType) => boolean;
 }
 
@@ -33,6 +35,7 @@ const useWorkspaceStore = create<WorkspaceState>()((set, get) => ({
     set({ ...get(), setting: workspaceSetting });
     return workspaceSetting;
   },
+  getSubscription: () => Subscription.fromPartial(get().profile.subscription || {}),
   checkFeatureAvailable: (feature: FeatureType): boolean => {
     return get().profile.subscription?.features.includes(feature) || false;
   },