improvement(usage-indicator): update query invalidation for usage to update in realtime (#2607)

* improvement(usage-indicator): update query invalidation for usage to update in realtime

* ack PR comments

Author: waleedlatif1

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts

@@ -252,7 +252,7 @@ export function useWand({
         })
 
         setTimeout(() => {
-          queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
+          queryClient.invalidateQueries({ queryKey: subscriptionKeys.all })
         }, 1000)
       } catch (error: any) {
         if (error.name === 'AbortError') {

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts

@@ -573,7 +573,7 @@ export function useWorkflowExecution() {
 
                 // Invalidate subscription queries to update usage
                 setTimeout(() => {
-                  queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
+                  queryClient.invalidateQueries({ queryKey: subscriptionKeys.all })
                 }, 1000)
 
                 safeEnqueue(encodeSSE({ event: 'final', data: result }))
@@ -646,7 +646,7 @@ export function useWorkflowExecution() {
 
           // Invalidate subscription queries to update usage
           setTimeout(() => {
-            queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
+            queryClient.invalidateQueries({ queryKey: subscriptionKeys.all })
           }, 1000)
         }
         return result

apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/components/cancel-subscription/cancel-subscription.tsx

@@ -165,7 +165,7 @@ export function CancelSubscription({ subscription, subscriptionData }: CancelSub
         logger.info('Subscription restored successfully', result)
       }
 
-      await queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
+      await queryClient.invalidateQueries({ queryKey: subscriptionKeys.all })
       if (activeOrgId) {
         await queryClient.invalidateQueries({ queryKey: organizationKeys.detail(activeOrgId) })
         await queryClient.invalidateQueries({ queryKey: organizationKeys.billing(activeOrgId) })

apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/usage-indicator/usage-indicator.tsx

@@ -199,7 +199,7 @@ export function UsageIndicator({ onClick }: UsageIndicatorProps) {
   useEffect(() => {
     const handleOperationConfirmed = () => {
       setTimeout(() => {
-        queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
+        queryClient.invalidateQueries({ queryKey: subscriptionKeys.all })
       }, 1000)
     }
     onOperationConfirmed(handleOperationConfirmed)

apps/sim/hooks/queries/subscription.ts

@@ -98,13 +98,13 @@ export function useUpdateUsageLimit() {
       return response.json()
     },
     onMutate: async ({ limit }) => {
-      await queryClient.cancelQueries({ queryKey: subscriptionKeys.user() })
-      await queryClient.cancelQueries({ queryKey: subscriptionKeys.usage() })
+      await queryClient.cancelQueries({ queryKey: subscriptionKeys.all })
 
-      const previousSubscriptionData = queryClient.getQueryData(subscriptionKeys.user())
+      const previousSubscriptionData = queryClient.getQueryData(subscriptionKeys.user(false))
+      const previousSubscriptionDataWithOrg = queryClient.getQueryData(subscriptionKeys.user(true))
       const previousUsageData = queryClient.getQueryData(subscriptionKeys.usage())
 
-      queryClient.setQueryData(subscriptionKeys.user(), (old: any) => {
+      const updateSubscriptionData = (old: any) => {
         if (!old) return old
         const currentUsage = old.data?.usage?.current || 0
         const newPercentUsed = limit > 0 ? (currentUsage / limit) * 100 : 0
@@ -120,7 +120,10 @@ export function useUpdateUsageLimit() {
             },
           },
         }
-      })
+      }
+
+      queryClient.setQueryData(subscriptionKeys.user(false), updateSubscriptionData)
+      queryClient.setQueryData(subscriptionKeys.user(true), updateSubscriptionData)
 
       queryClient.setQueryData(subscriptionKeys.usage(), (old: any) => {
         if (!old) return old
@@ -133,19 +136,24 @@ export function useUpdateUsageLimit() {
         }
       })
 
-      return { previousSubscriptionData, previousUsageData }
+      return { previousSubscriptionData, previousSubscriptionDataWithOrg, previousUsageData }
     },
     onError: (_err, _variables, context) => {
       if (context?.previousSubscriptionData) {
-        queryClient.setQueryData(subscriptionKeys.user(), context.previousSubscriptionData)
+        queryClient.setQueryData(subscriptionKeys.user(false), context.previousSubscriptionData)
+      }
+      if (context?.previousSubscriptionDataWithOrg) {
+        queryClient.setQueryData(
+          subscriptionKeys.user(true),
+          context.previousSubscriptionDataWithOrg
+        )
       }
       if (context?.previousUsageData) {
         queryClient.setQueryData(subscriptionKeys.usage(), context.previousUsageData)
       }
     },
     onSettled: () => {
-      queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
-      queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() })
+      queryClient.invalidateQueries({ queryKey: subscriptionKeys.all })
     },
   })
 }

apps/sim/stores/panel/copilot/store.ts

@@ -2661,7 +2661,7 @@ export const useCopilotStore = create<CopilotStore>()(
         // Invalidate subscription queries to update usage
         setTimeout(() => {
           const queryClient = getQueryClient()
-          queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
+          queryClient.invalidateQueries({ queryKey: subscriptionKeys.all })
         }, 1000)
       } finally {
         clearTimeout(timeoutId)