From 2e6f30322fe0c9d9aff2559eaaefc3d711168f26 Mon Sep 17 00:00:00 2001 From: waleed Date: Sat, 27 Dec 2025 14:44:43 -0800 Subject: [PATCH 1/2] improvement(usage-indicator): update query invalidation for usage to update in realtime --- .../workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts | 2 +- .../w/[workflowId]/hooks/use-workflow-execution.ts | 4 ++-- .../components/cancel-subscription/cancel-subscription.tsx | 2 +- .../sidebar/components/usage-indicator/usage-indicator.tsx | 2 +- apps/sim/hooks/queries/subscription.ts | 3 +-- apps/sim/stores/panel/copilot/store.ts | 2 +- 6 files changed, 7 insertions(+), 8 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts index 2f6d9bc47b..e1ed0500a2 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts +++ b/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') { diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts index f2d9cc2f76..61af9f0624 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts +++ b/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 diff --git a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/components/cancel-subscription/cancel-subscription.tsx b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/components/cancel-subscription/cancel-subscription.tsx index c778f06603..2461f3e739 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/components/cancel-subscription/cancel-subscription.tsx +++ b/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) }) diff --git a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/usage-indicator/usage-indicator.tsx b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/usage-indicator/usage-indicator.tsx index dd1df673ae..97f4dc8c36 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/usage-indicator/usage-indicator.tsx +++ b/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) diff --git a/apps/sim/hooks/queries/subscription.ts b/apps/sim/hooks/queries/subscription.ts index 750d5f5251..fd91b48464 100644 --- a/apps/sim/hooks/queries/subscription.ts +++ b/apps/sim/hooks/queries/subscription.ts @@ -144,8 +144,7 @@ export function useUpdateUsageLimit() { } }, onSettled: () => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() }) - queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }) + queryClient.invalidateQueries({ queryKey: subscriptionKeys.all }) }, }) } diff --git a/apps/sim/stores/panel/copilot/store.ts b/apps/sim/stores/panel/copilot/store.ts index f963983ce0..8a5a634af2 100644 --- a/apps/sim/stores/panel/copilot/store.ts +++ b/apps/sim/stores/panel/copilot/store.ts @@ -2661,7 +2661,7 @@ export const useCopilotStore = create()( // Invalidate subscription queries to update usage setTimeout(() => { const queryClient = getQueryClient() - queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() }) + queryClient.invalidateQueries({ queryKey: subscriptionKeys.all }) }, 1000) } finally { clearTimeout(timeoutId) From 32e2e6199cb2ce87ec4f4c34ae1a3ef8298761f8 Mon Sep 17 00:00:00 2001 From: waleed Date: Sat, 27 Dec 2025 15:11:37 -0800 Subject: [PATCH 2/2] ack PR comments --- apps/sim/hooks/queries/subscription.ts | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/apps/sim/hooks/queries/subscription.ts b/apps/sim/hooks/queries/subscription.ts index fd91b48464..89ded91231 100644 --- a/apps/sim/hooks/queries/subscription.ts +++ b/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,11 +136,17 @@ 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)